#!/usr/bin/python

import sys, os, getopt, libvirt

VERBOSE_LOG = False

def dlog(msg):
    if VERBOSE_LOG:
        print("debug: %s" % msg)

class Config:
    @staticmethod
    def rbd_to_file(uuid, rbd_path = ""):
        return "/vz/vmprivate/%s/%s" % (uuid, rbd_path.replace('/', '_'))

    @staticmethod
    def file_to_file(uuid, file_path = ""):
        return "/vz/vmprivate/%s/%s" % (uuid, file_path.strip("/").replace('/', '_'))

    @staticmethod
    def img_path_format(uuid, path = ""):
        if "rbd" in path:
            return Config.rbd_to_file(uuid, path.split(":")[1])
        else:
            return Config.file_to_file(uuid, path)

    def copy_vnc_addr(self, uuid, addr):
        if (addr == u'127.0.0.1' or addr == u'0.0.0.0'):
            return addr
        else:
            return self.opts.get("vzhost", "").split("@")[-1]

    def __init__(self, opts):
        self.opts  = opts
        self.convert = {
            "disk": {
                "network": {
                    "device": "disk",
                    "change-head-attrs": {
                        "type": {"val": "file"},
                    },
                    "items": {
                        "source": {
                            "type": "replace",
                            "data": {
                                "file": {"format": self.rbd_to_file},
                                "startupPolicy": {"val": "optional"},
                            },
                        },
                        "driver": {
                            "type": "change",
                            "data": {
                                "type" : {"val": "qcow2",},
                                "cache": {"val": "none",},
                                "detect_zeroes": {"val": "unmap"},
                            },
                        },
                        "auth": {
                            "type": "remove",
                        },
                    },
                },
                "file": {
                    "device": "disk",
                    "items": {
                         "source": {
                            "type": "change",
                            "data": {
                                "file": {"format": self.file_to_file},
                            },
                         },
                        "driver": {
                            "type": "change",
                            "data": {
                                "detect_zeroes": {"val": "unmap"},
                            },
                        },
                    },
                },
            },
            "interface": {
                "network": {
                    "device": "",
                    "change-head-attrs": {
                        "type": {"val": "bridge" if opts.get("netsource", "") else ""},
                    },
                     "items": {
                        "source": {
                            "type": "replace",
                            "data": {
                                "bridge": {"val": opts.get("netsource", "")},
                            },
                        },
                     },
                },
                "bridge": {
                    "device": "",
                    "items": {
                       "source": {
                           "type": "replace",
                           "data": {
                               "bridge": {"val": opts.get("netsource", "")},
                           },
                       },
                    },
                },
            },
            "graphics": {
                "vnc": {
                    "device": "",
                    "change-head-attrs": {
                        "listen": {"format": self.copy_vnc_addr},
                    },
                    "items": {
                        "listen": {
                            "type": "change",
                            "data": {
                                "address": {"format": self.copy_vnc_addr},
                            },
                        },
                    },
                },
            },
        }

class Convert:
    from xml.dom import minidom

    def __init__(self, config):
        self.config = config

    def __val_format(self, val, args):
        return val["format"](*args) if val.get("format", False) else val["val"]

    def __item_format(self, item, args):
        build = dict(item["data"])
        for name, value in item["data"].iteritems():
            if value.get("val") == "":
                return
            build[name] = self.__val_format(value, args)
        return build

    def __get_uuid(self, dom_xml):
        return iter(dom_xml.getElementsByTagName("uuid")).next().firstChild.nodeValue

    def xml(self, dom_xml_string):
        def create_node(root, name, args):
            if not args:
                return
            node = root.createElement(name)
            for name, key in args.iteritems():
                attr = root.createAttribute(name)
                attr.nodeValue = key
                node.setAttributeNode(attr)
            return node

        def replace_node(node, new, old):
            node.insertBefore(new, old)
            node.removeChild(old)

        def to_xml(root):
            import StringIO

            writer = StringIO.StringIO()
            for node in root.childNodes:
                node.writexml(writer, "", "", "")
            return writer.getvalue()

        def convert_by_config(root, config, subcfg):
            def change_items(node, attrs):
                for name, value in attrs.iteritems():
                    args = (self.__get_uuid(root), node.getAttribute(name))
                    new_value = self.__val_format(value, args)
                    if new_value:
                        node.setAttribute(name, new_value)

            def replace_items(sub, node, item):
                args = (self.__get_uuid(root), node.getAttribute("name"))
                new_node = create_node(root, node.nodeName, self.__item_format(item, args))
                if new_node:
                    replace_node(sub, new_node, node)

            cmd = {
                "replace": replace_items,
                "change": lambda _, node, item: change_items(node, item["data"]),
                "remove": lambda sub, node, _: sub.removeChild(node),
            }
            for sub in root.getElementsByTagName(subcfg):
                cfg = config[subcfg].get(sub.getAttribute("type"), False)
                if not cfg or cfg["device"] != sub.getAttribute("device"):
                    continue
                change_items(sub, cfg.get("change-head-attrs", {}))
                for node in sub.childNodes:
                    item = cfg["items"].get(node.nodeName, False)
                    if item:
                        cmd[item["type"]](sub, node, item)

        root = self.minidom.parseString(dom_xml_string)
        for subcfg in self.config:
            convert_by_config(root, self.config, subcfg)

        xml_string_converted = to_xml(root)
        root.unlink()
        return xml_string_converted

def qmp_monitor(dom, monitor_cmd):
    import libvirt_qemu, json
    return json.loads(libvirt_qemu.qemuMonitorCommand(dom, monitor_cmd, 0))["return"]

def prepare_storage_on_vz(dom, vzconn):
    vzmigrate_pool = """
<pool type="dir">
  <name>%s</name>
  <target>
    <path>%s</path>
  </target>
</pool>"""

    vzmigrate_vol = """
<volume>
  <name>%s</name>
  <capacity unit='bytes'>%d</capacity>
  <target>
    <format type='qcow2'/>
    <compat>1.1</compat>
  </target>
</volume>"""

    VIR_STORAGE_POOL_CREATE_WITH_BUILD = 1
    uuid = dom.UUIDString()
    try:
        vzpool = vzconn.storagePoolLookupByName(uuid)
        vzpool.destroy()
    except:
        dlog("pool %s already exists" % uuid)
    xml_pool = vzmigrate_pool % (uuid, os.path.dirname(Config.img_path_format(uuid)))
    vzpool = vzconn.storagePoolCreateXML(xml_pool, VIR_STORAGE_POOL_CREATE_WITH_BUILD)
    if vzpool == None:
        print("Failed to create StoragePool object")
        sys.exit(1)
    dlog("vzpool created %s" % os.path.dirname(Config.img_path_format(uuid)))

    disks = qmp_monitor(dom, "{ \"execute\": \"query-block\" }")
    vzvols = []
    for disk in filter(lambda disk: not disk["removable"], disks):
        img = disk["inserted"]["image"]
        img_path = Config.img_path_format(uuid, img["filename"]) #todo: keep legacy format
        vzvol = vzpool.createXML(vzmigrate_vol % (os.path.basename(img_path), img["virtual-size"]))
        vzvols.append(vzvol)
        if vzvol == None:
            print('Failed to create a StorageVol objects')
            vzpool.destroy()
            sys.exit(1)
        dlog("volume created %s" % img_path)
        #vzvol.delete(0) #logically remove the storage volume
    return vzpool, vzvols

def main(argv):
    def help():
        print("usage:\nrhel_to_vz.py [-v] [--vzhost=addr] [--vmname=name] [--netsource=ifname] [--migration-port=port1] [--mirror-port=port2]\n")
        print("example:\nrhel_to_vz.py --vzhost=user@vzhost.com --vmname=VM123\n")
        print("note:\nMake sure that migration and mirroring ports are opened in the firewall and are not used.")
        print("To make work port specifying the last version of libvirt-python package is required.")
        print("Details: https://www.redhat.com/archives/libvir-list/2018-February/msg00351.html\n")
        print(" --vzhost\t\t Destination address")
        print(" --vmname\t\t Target vm name")
        print(" --netsource\t\t Set new network source [optional]")
        print(" --migration-port\t Set migration port [optional] (default is 49152)")
        print(" --mirror-port\t\t Set block mirroring port [optional] (default is 49153)")
        print(" -v, --verbose\t\t Enable debug messages")
        print(" -h, --help\t\t Display this help and exit")
    def usage():
        print("wrong parameters")
        help()
        sys.exit(1)

    try:
        opt_list, args = getopt.getopt(sys.argv[1:], "hv",
                                   ["help", "verbose", "vzhost=", "vmname=",
                                    "netsource=", "migration-port=",
                                    "mirror-port="])

    except getopt.GetoptError as err:
        print str(err)
        usage()

    opts = {}
    for arg, val in opt_list:
        if arg in ("-v", "--verbose"):
            global VERBOSE_LOG
            VERBOSE_LOG = True
        elif arg in ("-h", "--help"):
            help()
            sys.exit(0)
        elif arg in ("-s", "--vzhost"):
            opts["vzhost"] = val
        elif arg in ("-m", "--vmname"):
            opts["vmname"] = val
        elif arg in ("-n", "--netsource"):
            opts["netsource"] = val
        elif arg == "--migration-port":
            opts["migration-port"] = val
        elif arg == "--mirror-port":
            opts["mirror-port"] = val
        else:
            assert False, "unhandled option %s" % arg
    dlog(opt_list)
    if not opts.get("vzhost", False) or not opts.get("vmname", False):
        usage()

    conn = libvirt.open('qemu:///system')
    if conn == None:
        print("Failed to open connection to qemu:///system")
        sys.exit(1)

    dom = conn.lookupByName(opts["vmname"])
    if dom == None:
        print("Failed to get the domain object, domname = %s" % opts["vmname"])
        sys.exit(1)

    orig_xml = dom.XMLDesc(2)
    vz_xml = Convert(Config(opts).convert).xml(orig_xml)
    dlog(vz_xml)
    dlog("domain xml converted, xml size %d" % len(vz_xml))

    vzconn = libvirt.openAuth("qemu+ssh://%s/system" % opts["vzhost"], [[], None, None], 0)
    if vzconn == None:
        print("Failed to open connection to qemu+ssh://%s/system" % opts["vzhost"])
        sys.exit(1)

    vzpool, vzvols = prepare_storage_on_vz(dom, vzconn)
    dlog("storage prepared")

    params = {libvirt.VIR_MIGRATE_PARAM_DEST_XML: vz_xml.encode('UTF-8')}
    if opts.get("migration-port", False):
        params[libvirt.VIR_MIGRATE_PARAM_URI] = "tcp://" + opts["vzhost"] + ":" + opts["migration-port"]
    if opts.get("mirror-port", False):
        params[libvirt.VIR_MIGRATE_PARAM_DISKS_PORT] = int(opts["mirror-port"])

    flags = (libvirt.VIR_MIGRATE_LIVE | libvirt.VIR_MIGRATE_PERSIST_DEST | libvirt.VIR_MIGRATE_NON_SHARED_DISK |
             libvirt.VIR_MIGRATE_AUTO_CONVERGE | libvirt.VIR_MIGRATE_COMPRESSED)
    print("migration %s to vz.." % opts["vmname"])
    qmp_monitor(dom, "{ \"execute\": \"migrate_set_downtime\", \"arguments\": { \"value\": 3 } }")
    dom.undefine()
    vzdom = None
    try:
        dom = conn.defineXML(vz_xml)
        vzdom = dom.migrate3(vzconn, params, flags)
    except:
        for vzvol in vzvols:
            vzvol.wipe(0)
            vzvol.delete(0)
    finally:
        conn.defineXML(orig_xml)
    if vzdom == None:
        print("Migration %s to %s failed" % (opts["vmname"], opts["vzhost"]))
        sys.exit(1)
    dlog("%s migrated" % opts["vmname"])

    print("migration finished with success.")

if __name__ == "__main__":
    main(sys.argv)