ova_import/ova_import

#!/usr/bin/env python3

# Copyright (C) 2016-2018 Bernhard Ehlers
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
ova_import generates gns3a file to import ova appliance into GNS3.

usage: ova_import [-h] [-d DIR] [-k] ova

positional arguments:
  ova                   .ova appliance

optional arguments:
  -h, --help            show this help message and exit
  -d DIR, --dir DIR     directory for storing appliance and images
  -k, --keepcompressed  keep images compressed
"""

import os
import sys
import argparse
import hashlib
import json
import struct
import tarfile
import zlib
import xml.etree.ElementTree as ElementTree
from collections import defaultdict, OrderedDict


def strip_namespace(tag):
    """ strip namespace from tag """
    if '}' in tag:
        tag = tag.split('}', 1)[1]      # strip all namespaces
    return tag


def etree_to_dict(etree):
    """
    convert ElementTree to dictionary
    see K3---rnc's answer of "Converting xml to dictionary using ElementTree"
    http://stackoverflow.com/questions/7684333/converting-xml-to-dictionary-using-elementtree#answer-10076823
    variable names changed to make it easier understandable
    """
    tag = strip_namespace(etree.tag)
    node = {tag: {} if etree.attrib else None}
    children = list(etree)
    if children:
        all_child_nodes = defaultdict(list)
        for child_node in map(etree_to_dict, children):
            for key, val in child_node.items():
                all_child_nodes[key].append(val)
        node = {tag: {key:val[0] if len(val) == 1 else val
                      for key, val in all_child_nodes.items()}}
    if etree.attrib:
        node[tag].update((strip_namespace(key), val)
                         for key, val in etree.attrib.items())
    if etree.text:
        text = etree.text.strip()
        if children or etree.attrib:
            if text:
                node[tag]['#text'] = text
        else:
            node[tag] = text
    return node


def create_gns3a(fname, dname, file_info, ovf_data):
    """ create GNS3 appliance file """
    fbase = os.path.basename(fname)
    gns3a = OrderedDict()
    ovf = etree_to_dict(ElementTree.fromstring(ovf_data))['Envelope']

    # base informations
    try:
        vm_name = ovf['VirtualSystem']['Name']
    except KeyError:
        vm_name = ovf['VirtualSystem']['id']
    gns3a["name"]             = vm_name
    gns3a["category"]         = "guest"
    gns3a["description"]      = "ova import of " + fbase
    gns3a["vendor_name"]      = "unknown"
    gns3a["vendor_url"]       = "http://www.example.com"
    gns3a["product_name"]     = vm_name
    gns3a["registry_version"] = 3
    gns3a["status"]           = "experimental"
    gns3a["maintainer"]       = "GNS3 Team"
    gns3a["maintainer_email"] = "developers@gns3.net"

    # qemu
    cores = 0
    eth_adapters = 0
    ram = 0
    for item in ovf['VirtualSystem']['VirtualHardwareSection']['Item']:
        if item['ResourceType'] == '3':
            cores = int(item['VirtualQuantity'])
        elif item['ResourceType'] == '4':
            ram = int(item['VirtualQuantity'])
        elif item['ResourceType'] == '10':
            eth_adapters += 1
    cores = max(cores, 1)
    eth_adapters = max(eth_adapters, 1)
    if ram == 0:
        ram = 256
    try:
        vm_os = ovf['VirtualSystem']['OperatingSystemSection']['osType']
    except KeyError:
        vm_os = ovf['VirtualSystem']['OperatingSystemSection']['OSType']['#text']

    qemu = OrderedDict()
    qemu["adapter_type"] = "e1000"
    qemu["adapters"] = eth_adapters
    qemu["ram"] = ram
    qemu["arch"] = "i386"
    if "64" in vm_os:
        qemu["arch"] = "x86_64"
    qemu["console_type"] = "telnet"
    qemu["kvm"] = "allow"
    if cores > 1:
        qemu["options"] = "-smp " + str(cores)
    gns3a["qemu"] = qemu

    # images
    images = []
    for name in sorted(file_info.keys()):
        image = OrderedDict()
        image["filename"] = name
        image["version"]  = "0.0"
        image["md5sum"]   = file_info[name]["md5"]
        image["filesize"] = file_info[name]["len"]
        images.append(image)
    gns3a["images"] = images

    # versions
    images = OrderedDict()
    cdrom = None
    disk_id = 0
    ovf_ref_file = ovf['References']['File']
    if not isinstance(ovf_ref_file, (list, tuple)):
        ovf_ref_file = [ovf_ref_file]
    for image in ovf_ref_file:
        img = image['href']
        if img.endswith(".iso"):
            if cdrom is None:
                cdrom = img
        else:
            images["hd" + chr(ord("a")+disk_id) + "_disk_image"] = img
            disk_id += 1
    if cdrom is not None:
        images["cdrom_image"] = cdrom
    gns3a["versions"] = [OrderedDict([("name", "0.0"), ("images", images)])]

    # write to file
    ofile = os.path.join(dname, os.path.splitext(fbase)[0] + ".gns3a")
    with open(ofile, "w") as f_out:
        json.dump(gns3a, f_out, indent=4, separators=(',', ': '))
        f_out.write("\n")


SECTOR_SIZE = 512

def div_ceil(dividend, divisor):
    """ integer division, round up """
    return -(-dividend // divisor)


def copy_raw(f_in, f_out, initial_data=None):
    """ copy input file to output file, return file info """
    f_len = 0
    f_md5 = hashlib.md5()

    if initial_data:
        f_out.write(initial_data)
        f_len += len(initial_data)
        f_md5.update(initial_data)

    while True:
        block = f_in.read(64 * 1024)
        if not block:
            break
        f_out.write(block)
        f_len += len(block)
        f_md5.update(block)

    return {'len': f_len, 'md5': f_md5.hexdigest()}


def copy_vmdk(f_in, f_out, out_name):
    """ uncompress VMDK file, return file info """

    header = f_in.read(SECTOR_SIZE)
    if len(header) != SECTOR_SIZE or header[0:4] != b'KDMV':
        # unknown file type: extract unchanged
        return copy_raw(f_in, f_out, header)

    # decode header
    (magic, _, flags, capacity, grain_size, descriptor_offset, descriptor_size,
     num_gte_per_gt, _, _, overhead, _, single_eol_char, non_eol_char,
     double_eol_char1, double_eol_char2, compress_algorithm) = \
        struct.unpack_from('<LLLQQQQLQQQBccccH', header)
    if flags & 0x30000 != 0x30000 or compress_algorithm != 1:
        # uncompressed VMDK file: extract unchanged
        return copy_raw(f_in, f_out, header)

    # grain directory / table
    gt_size = div_ceil(num_gte_per_gt * 4, SECTOR_SIZE)
    gt_count = div_ceil(capacity, grain_size * num_gte_per_gt)
    gd_size = div_ceil(gt_count * 4, SECTOR_SIZE)
    rgd_offset = 8
    gd_offset = rgd_offset + gd_size + gt_count*gt_size
    out_overhead = div_ceil(gd_offset + gd_size + gt_count*gt_size, 128) * 128

    # read/modify descriptor
    if descriptor_offset > 0 and descriptor_size > 0:
        skip_len = (descriptor_offset-1) * SECTOR_SIZE
        if skip_len > 0 and len(f_in.read(skip_len)) != skip_len:
            raise ValueError('Premature EOF')
        descriptor = f_in.read(descriptor_size * SECTOR_SIZE)
        if len(descriptor) != descriptor_size * SECTOR_SIZE:
            raise ValueError('Premature EOF')
        in_offset = descriptor_offset + descriptor_size
        descriptor = descriptor.rstrip(b'\0')
        descriptor = descriptor.replace(double_eol_char1 + double_eol_char2,
                                        single_eol_char)
        if descriptor[-1:] != single_eol_char:
            descriptor += single_eol_char
        while non_eol_char + single_eol_char in descriptor:
            descriptor = descriptor.replace(non_eol_char + single_eol_char,
                                            single_eol_char)
        descriptor = descriptor.replace(b'"streamOptimized"',
                                        b'"monolithicSparse"')
        off_extend = descriptor.find(single_eol_char + b'RDONLY ')
        if off_extend < 0:
            raise ValueError('No extent defined in descriptor')
        off_extend += 1
        off_extend_end = descriptor.find(single_eol_char, off_extend)
        if off_extend_end < 0:
            off_extend_end = len(descriptor)
        if descriptor.find(b'RDONLY ', off_extend_end) >= 0:
            raise ValueError('Multiple extents defined in descriptor')
        descriptor = descriptor[0:off_extend] + \
                     'RW {:d} SPARSE "{}"'.format(capacity, out_name).encode('utf-8') + \
                     descriptor[off_extend_end:]
        out_desc_offset = 1
        out_desc_size = rgd_offset - 1
        if len(descriptor) > out_desc_size * SECTOR_SIZE:
            raise ValueError('Descriptor too big')
    else:
        in_offset = 1
        out_desc_offset = 0
        out_desc_size = 0

    # write new header
    header = struct.pack('<LLLQQQQLQQQBccccH', \
        magic, 1, 3, capacity, grain_size, out_desc_offset, out_desc_size, \
        num_gte_per_gt, rgd_offset, gd_offset, out_overhead, 0, \
        single_eol_char, non_eol_char, double_eol_char1, double_eol_char2, 0)
    header += b'\0' * (SECTOR_SIZE - len(header))
    f_out.write(header)

    # write descriptor
    if out_desc_offset > 0:
        descriptor += b'\0' * (out_desc_size*SECTOR_SIZE - len(descriptor))
        f_out.write(descriptor)

    # skip to data
    skip_len = (overhead - in_offset) * SECTOR_SIZE
    if skip_len > 0 and len(f_in.read(skip_len)) != skip_len:
        raise ValueError('Premature EOF')
    in_offset = overhead
    out_offset = out_overhead
    f_out.seek(out_offset * SECTOR_SIZE, 0)

    # read, uncompress and write data
    grain_directory = None
    grain_table = []
    grain_trans_init = {0:0, 1:1}
    grain_translation = grain_trans_init.copy()
    gt_translation = {0:-1}
    max_datasize = grain_size * SECTOR_SIZE
    # add worst case compression overhead
    max_datasize += div_ceil(max_datasize, 8) + div_ceil(max_datasize, 64) + 11
    while True:
        header = f_in.read(SECTOR_SIZE)
        if len(header) < SECTOR_SIZE:
            raise ValueError('Premature EOF')
        (val, size, marker) = struct.unpack_from('<QLL', header)

        if size > 0:
            # Compressed Grain
            if size > max_datasize:
                raise ValueError('Bad data block @{:08X}'.format(in_offset))
            prev_offset = in_offset
            grain_translation[in_offset] = out_offset
            data = header[12:]
            if size > (SECTOR_SIZE - 12):
                data += f_in.read(size - (SECTOR_SIZE - 12))
                if len(data) != size:
                    raise ValueError('Premature EOF')
                padding = -(size + 12) % SECTOR_SIZE
                if padding > 0 and len(f_in.read(padding)) != padding:
                    raise ValueError('Premature EOF')
                in_offset += div_ceil(size + 12, SECTOR_SIZE)
            else:
                data = data[0:size]
                in_offset += 1
            try:
                data = zlib.decompress(data)
            except zlib.error:
                raise ValueError('Bad data block @{:08X}'.format(prev_offset))
            if len(data) != grain_size*SECTOR_SIZE:
                raise ValueError('Data block @{:08X} has wrong size {:d}'.format(prev_offset, len(data)))
            f_out.write(data)
            out_offset += grain_size

        elif marker == 0:
            # End-of-Stream
            if f_in.read(SECTOR_SIZE):
                raise ValueError('No EOF after end of stream @{:08X}'.format(in_offset))
            break

        elif marker == 1:
            # Grain Table
            if len(grain_table) >= gt_count:
                raise ValueError('Too many grain tables @{:08X}'.format(in_offset))
            if val != gt_size:
                raise ValueError('Grain table @{:08X} has wrong size {:d}'.format(in_offset, val * SECTOR_SIZE))
            gt_translation[in_offset+1] = len(grain_table)
            data = f_in.read(val * SECTOR_SIZE)
            if len(data) != val * SECTOR_SIZE:
                raise ValueError('Premature EOF')
            data = struct.unpack_from("<{:d}L".format(num_gte_per_gt), data)
            # apply new offsets in grain tables
            try:
                data = [grain_translation[x] for x in data]
            except KeyError:
                raise ValueError('Grain table uses invalid data references')
            grain_table.append(data)
            grain_translation = grain_trans_init.copy()
            in_offset += 1 + val

        elif marker == 2:
            # Grain Directory
            if grain_directory is not None:
                raise ValueError('Too many grain directories @{:08X}'.format(in_offset))
            if val != gd_size:
                raise ValueError('Grain directory @{:08X} has wrong size {:d}'.format(in_offset, val * SECTOR_SIZE))
            data = f_in.read(val * SECTOR_SIZE)
            if len(data) != val * SECTOR_SIZE:
                raise ValueError('Premature EOF')
            data = struct.unpack_from("<{:d}L".format(gt_count), data)
            grain_directory = data
            in_offset += 1 + val

        else:
            # ignore footer or unknown marker
            if val > 0 and len(f_in.read(val*SECTOR_SIZE)) != val*SECTOR_SIZE:
                raise ValueError('Premature EOF')
            in_offset += 1 + val

    # set table index in grain directory
    if grain_directory is None:
        raise ValueError('No grain directory')
    try:
        grain_directory = [gt_translation[x] for x in grain_directory]
    except KeyError:
        raise ValueError('Grain directory uses invalid table references')

    # grain directory: sort grain tables, update table index to sector offset
    gt_empty = [0] * num_gte_per_gt
    unsorted_grain_table = grain_table
    grain_table = []
    gt_offset = rgd_offset + gd_size
    for index in range(0, gt_count):
        if grain_directory[index] < 0:
            grain_table.append(gt_empty)
        else:
            grain_table.append(unsorted_grain_table[grain_directory[index]])
        grain_directory[index] = gt_offset
        gt_offset += gt_size

    # save first/redundant grain descriptor
    f_out.seek(rgd_offset * SECTOR_SIZE, 0)
    data = struct.pack("<{:d}L".format(gt_count), *grain_directory)
    data += b'\0' * (gd_size*SECTOR_SIZE - len(data))
    f_out.write(data)
    for table in grain_table:
        data = struct.pack("<{:d}L".format(num_gte_per_gt), *table)
        data += b'\0' * (gt_size*SECTOR_SIZE - len(data))
        f_out.write(data)

    # save second grain descriptor
    second_directory = [x - rgd_offset + gd_offset for x in grain_directory]
    data = struct.pack("<{:d}L".format(gt_count), *second_directory)
    data += b'\0' * (gd_size*SECTOR_SIZE - len(data))
    f_out.write(data)
    for table in grain_table:
        data = struct.pack("<{:d}L".format(num_gte_per_gt), *table)
        data += b'\0' * (gt_size*SECTOR_SIZE - len(data))
        f_out.write(data)

    # get file information
    f_out.seek(0, 0)
    f_len = 0
    f_md5 = hashlib.md5()
    while True:
        block = f_out.read(64 * 1024)
        if not block:
            break
        f_len += len(block)
        f_md5.update(block)

    return {'len': f_len, 'md5': f_md5.hexdigest()}


def main(argv):
    """ main """

    # parse command line
    parser = argparse.ArgumentParser(description='%(prog)s generates gns3a file to import ova appliance into GNS3.')
    parser.add_argument('-d', '--dir', default='',
                        help='directory for storing appliance and images')
    parser.add_argument('-k', '--keepcompressed', action='store_true',
                        help='keep images compressed')
    parser.add_argument('ova', help='.ova appliance')
    args = parser.parse_args(argv[1:])

    fname = args.ova
    dname = args.dir

    if dname != '' and not os.path.isdir(dname):
        sys.exit("Directory '{}' doesn't exist.".format(dname))

    # open ovf file
    try:
        tar = tarfile.open(fname)
    except (IOError, tarfile.TarError) as err:
        sys.exit("Error reading ova file: {}".format(err))

    # get files from ova
    file_info = {}
    ovf_data = None

    for tarinfo in tar:
        name = tarinfo.name.split('/')[-1]
        if not tarinfo.isfile():            # ova uses only regular files
            pass
        elif name.endswith(".ovf"):         # ovf file
            f_in = tar.extractfile(tarinfo)
            ovf_data = f_in.read()
            f_in.close()
        elif name.endswith(".mf"):          # ignore manifest file
            pass
        elif name.endswith(".cert"):        # ignore certificate
            pass
        else:                               # save image file
            print("Extracting image {}...".format(name))
            out_name = os.path.join(dname, name)
            f_in = tar.extractfile(tarinfo)
            f_out = open(out_name, 'w+b')
            try:
                if not args.keepcompressed and name.endswith(".vmdk"):
                    # uncompress VMDK file
                    file_info[name] = copy_vmdk(f_in, f_out, name)
                else:
                    file_info[name] = copy_raw(f_in, f_out)
            except ValueError as err:
                sys.exit("{}: {}".format(name, err))
            f_in.close()
            f_out.close()
            os.utime(out_name, (tarinfo.mtime, tarinfo.mtime))

    tar.close()

    if ovf_data is None:
        sys.exit("No ovf information in ova.")

    print("Creating appliance {}..."
          .format(os.path.splitext(os.path.basename(fname))[0] + ".gns3a"))
    create_gns3a(fname, dname, file_info, ovf_data)


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