#!/usr/bin/env python
# encoding: utf-8

# Copyright (C) 2011 Red Hat, Inc.
#
# 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; version 2 of the License.
#
# 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
# MA  02110-1301, USA.  A copy of the GNU General Public License is
# also available at http://www.gnu.org/copyleft/gpl.html.
#
# First cut module to convert raw disk images into stream-optimized VMDK files
#
# See the "Specification Document" referenced in Wikipedia for more details:
#
# http://en.wikipedia.org/wiki/VMDK
#
# Divergence from spec noted below
#
# The stream-optimized format is required for importing images via the vSphere
# SOAP API

import struct
import sys
import os
import math
import string
import zlib

class VMDKStreamException(Exception):
    def __init__(self, msg):
        self.msg = msg
    def __str__(self):
        return self.msg

# Header Constants
MAGIC_NUMBER = 0x564D444B # 'V' 'M' 'D' 'K'

# Marker Constants
MARKER_EOS = 0 # end of stream
MARKER_GT = 1 # grain table
MARKER_GD = 2 # grain directory
MARKER_FOOTER = 3 # footer (repeat of header with final info)

# Other Constants
SECTOR_SIZE = 512

# Descriptor Template
image_descriptor_template='''# Description file created by VMDK stream converter
version=1
# Believe this is random
CID=7e5b80a7
# Indicates no parent
parentCID=ffffffff
createType="streamOptimized"

# Extent description
RDONLY #SECTORS# SPARSE "call-me-stream.vmdk"

# The Disk Data Base
#DDB

ddb.adapterType = "lsilogic"
# #SECTORS# / 63 / 255 rounded up
ddb.geometry.cylinders = "#CYLINDERS#"
ddb.geometry.heads = "255"
ddb.geometry.sectors = "63"
# Believe this is random
ddb.longContentID = "8f15b3d0009d9a3f456ff7b28d324d2a"
ddb.virtualHWVersion = "7"'''


def create_sparse_header(inFileSectors, descriptorSize,
                         gdOffset = 0xFFFFFFFFFFFFFFFF):
    # While theoretically variable we set these based on current VMWare
    # convention
    grainSize = 128
    numGTEsPerGT = 512
    overHead = 128
    formatVersion = 3 # NOTE: Conflicts with VMWare docs - determined by trial/error

    descriptorOffset = 1

    # The following are always fixed in the "stream-optimized" format we are
    # creating
    compressAlgorithm = 1
    flags = 0x30001
    rgdOffset = 0

    # We are building from scratch so an unclean shutdown is not possible
    uncleanShutdown = 0

    # Build the struct
    header_list = [ MAGIC_NUMBER, formatVersion, flags, inFileSectors,
                    grainSize, descriptorOffset, descriptorSize, numGTEsPerGT,
                    rgdOffset, gdOffset, overHead, uncleanShutdown,
                    '\n', ' ', '\r', '\n', compressAlgorithm ]
    for i in range(433):
	header_list.append(0)
    header_struct = "=IIIQQQQIQQQBccccH433B"
    return struct.pack(header_struct, *header_list)

def create_marker(numSectors, size, marker_type):
    marker_list = [ numSectors, size, marker_type ]
    for i in range(496):
	marker_list.append(0)
    marker_struct = "=QII496B"
    return struct.pack(marker_struct, *marker_list)

def create_grain_marker(location, size):
    # The grain marker is special in that the data follows immediately after it
    # without a pad
    return struct.pack("=QI", location, size)

def divro(num, den):
    # Divide always rounding up and returning an integer
    # Is there some nicer way to do this?
    return int(math.ceil((1.0*num)/(1.0*den)))

def pad_to_sector(stringlen):
    # create a pad that, when concated onto a string of stringlen
    # makes it an integer number of sectors
    # return pad and length of input_string + pad in sectors
    pad = ""
    if stringlen % SECTOR_SIZE:
        # This does need padding
        for i in range(SECTOR_SIZE - (stringlen % SECTOR_SIZE)):
            pad += '\0'
    finallen = (stringlen + len(pad))/SECTOR_SIZE
    return pad, finallen

def sector_pointer(file_object):
    # return file point in sectors
    # raise an exception if not sector aligned
    file_location = file_object.tell()
    if file_location % SECTOR_SIZE:
        raise VMDKStreamException("Asked for a sector pointer on a file whose r/w pointer is not sector aligned")
    else:
        return file_location / SECTOR_SIZE


def write_grain_table(outfile, grain_table, gtes_per_gt = 512):
    # Write grain_table to outfile including header
    # return the sector on which the table starts

    zero_grain_table = [ ]
    for i in range(gtes_per_gt):
        zero_grain_table.append(0)

    if grain_table == zero_grain_table:
        # We don't need to write this and can put zeros in the directory
        return 0
    else:
        grain_table_marker = create_marker(numSectors = (gtes_per_gt * 4) / SECTOR_SIZE,
                                           size = 0, marker_type = MARKER_GT)
        outfile.write(grain_table_marker)
        table_location = sector_pointer(outfile)
        outfile.write(struct.pack("%dI" % gtes_per_gt, *grain_table))
        return table_location

def debug_print(message):
    #print message
    pass

def convert_to_stream(infilename, outfilename):
    debug_print("DEBUG: opening %s to write to %s" % (infilename, outfilename))

    infileSize = os.path.getsize(infilename)
    infileSectors = divro(infileSize, 512)
    debug_print("DEBUG: input file is (%s) bytes - (%s) sectors long" % (infileSize, infileSectors))

    # Fixed by convention
    # TODO: Make variable here and in header fuction
    grainSectors=128
    totalGrains=divro(infileSectors, grainSectors)
    debug_print("DEBUG: total grains will be (%s)" % (totalGrains))

    # Fixed by convention
    # TODO: Make variable here and in header fuction
    numGTEsPerGT = 512
    totalGrainTables=divro(totalGrains, numGTEsPerGT)
    debug_print("DEBUG: total Grain Tables needed will be (%s)" % (totalGrainTables))

    grainDirectorySectors=divro(totalGrainTables*4, SECTOR_SIZE)
    debug_print("DEBUG: sectors in Grain Directory will be (%s)" % (grainDirectorySectors))

    grainDirectoryEntries=grainDirectorySectors*128
    debug_print("DEBUG: Number of entries in Grain Directory - (%s)" % (grainDirectoryEntries))

    infileCylinders=divro(infileSectors, (63*255))
    debug_print("DEBUG: Cylinders (%s)" % infileCylinders)

    # Populate descriptor
    tmpl = image_descriptor_template
    tmpl = string.replace(tmpl, "#SECTORS#", str(infileSectors))
    tmpl = string.replace(tmpl, "#CYLINDERS#", str(infileCylinders))
    image_descriptor = tmpl

    image_descriptor_pad, desc_sectors = pad_to_sector(len(image_descriptor))
    debug_print("DEBUG: Descriptor takes up (%s) sectors" % desc_sectors)
    image_descriptor += image_descriptor_pad

    image_header = create_sparse_header(inFileSectors = infileSectors,
                                        descriptorSize = desc_sectors)

    outfile = open(outfilename, "wb")
    outfile.write(image_header)
    outfile.write(image_descriptor)

    # Fixed by convention
    # TODO: Make variable here and in header function
    overHead = 128

    # Pad the output file to fill the overHead
    for i in range((overHead-sector_pointer(outfile)) * SECTOR_SIZE):
        outfile.write('\0')

    # grainDirectory - list of integers representing the global level 0 grain
    # directory
    grainDirectory = [ ]

    # currentGrainTable - list that can grow to numGTEsPerGT integers
    # representing the active grain table
    currentGrainTable = [ ]

    # For slightly more efficient comparison
    grainSize = grainSectors * SECTOR_SIZE
    zeroChunk = ""
    for i in range(grainSize):
        zeroChunk += '\0'

    # We are ready to start reading
    infile = open(infilename, "rb")

    try:
        inputSectorPointer = sector_pointer(infile)
        inChunk = infile.read(grainSize)
        while inChunk != "":
            if inChunk == zeroChunk:
                # All zeros - no need to create a grain - just mark zero in GTE
                currentGrainTable.append(0)
            else:
                # Create a compressed grain
                currentGrainTable.append(sector_pointer(outfile))
		compChunk = zlib.compress(inChunk)
                grain_marker = create_grain_marker(inputSectorPointer,
                                                   len(compChunk))
		grainPad, writeSectors = pad_to_sector(len(compChunk) + len(grain_marker))
		outfile.write(grain_marker)
		outfile.write(compChunk)
		outfile.write(grainPad)

            if len(currentGrainTable) == numGTEsPerGT:
                # Table is full
                table_location =  write_grain_table(outfile, currentGrainTable,
                                                    gtes_per_gt = numGTEsPerGT)
                # function does zero check so we don't have to
                grainDirectory.append(table_location)
                currentGrainTable = [ ]
            # do not update pointer unless we read a full grain last time
            # incomplete grain read indicates EOF and may result in non-sector
            # alignment
            if len(inChunk) == grainSize:
                inputSectorPointer = sector_pointer(infile)
            # read the next chunk
            inChunk = infile.read(grainSize)
    finally:
        # Write out the final grain table if needed
        if len(currentGrainTable):
            debug_print("Partial grain table present - padding and adding it to dir")
            for i in range(numGTEsPerGT-len(currentGrainTable)):
                currentGrainTable.append(0)
            table_location = write_grain_table(outfile, currentGrainTable,
                                               gtes_per_gt = numGTEsPerGT)
            grainDirectory.append(table_location)
        else:
            debug_print("Current grain table is empty so we need not write it out")

        # pad out grain directory then write it
        for i in range(grainDirectoryEntries - totalGrainTables):
            grainDirectory.append(0)
        grain_directory_marker = create_marker(grainDirectorySectors, 0,
                                               MARKER_GD)
        outfile.write(grain_directory_marker)
        gdLocation = sector_pointer(outfile)
        grainDirectoryStruct = "%dI" % grainDirectoryEntries
        debug_print("Grain directory length (%d)" % (len(grainDirectory)))
        debug_print("Grain directory: ")
        debug_print(grainDirectory)
        outfile.write(struct.pack(grainDirectoryStruct, *grainDirectory))

        # footer marker
        outfile.write(create_marker(1, 0, MARKER_FOOTER))

        # footer
        footer = create_sparse_header(inFileSectors = infileSectors,
                                      descriptorSize = desc_sectors,
                                      gdOffset = gdLocation)
        outfile.write(footer)

        # EOS marker
        outfile.write(create_marker(0, 0, MARKER_EOS))
        outfile.close()
        infile.close()

if __name__ == '__main__':
    convert_to_stream(sys.argv[1], sys.argv[2])
