#!/usr/bin/python
# 
# olpc-image-builder
# 
# Description: A python script to assist with the creation of partitioned
#              images for the OLPC XO-1.
# Documentation: http://wiki.laptop.org/go/OFW_NAND_FLASH_Updater
# Author:  Erik Garrison <erik@laptop.org>

import hashlib
import os


KiB_bytes = pow(2,10)
MiB_bytes = pow(KiB_bytes, 2)
GiB_bytes = pow(KiB_bytes, 3)

# TODO is rounding an issue here?
#      should we perhaps be doing
#    if size_spec[-3:] == "GiB":
#        return float(GiB_bytes) * int(size_spec[:-3])
# Hmm.
def human_size_to_bytes(size_spec):
    if size_spec[-3:] == "GiB":
        return GiB_bytes * int(size_spec[:-3])
    elif size_spec[-3:] == "MiB":
        return MiB_bytes * int(size_spec[:-3])
    elif size_spec[-3:] == "KiB":
        return KiB_bytes * int(size_spec[:-3])
    else: # assume size is in bytes
        return int(size_spec)

def dec2hex(n):
    """return the hexadecimal string representation of integer n"""
    return "%x" % n
 
def hex2dec(s):
    """return the integer value of a hexadecimal string s"""
    return int(s, 16)

# from: http://wiki.laptop.org/go/OFW_NAND_FLASH_Updater
# for each partition, do this....
"""
set-partition: 1
mark-pending: 0
eblock: 0 sha256 ff3227e680963cb0a5e24fdbdaf1aa5437c9ce400761f7c5a462d75f1f18f204
eblock: 1 sha256 289a9e81513229f14d7ee6c0d9eb588ffd7c7e87e146da42b2102a0cc2f9598f
...
eblock: 37 sha256 02eb0c9c03c10bf9cc3a590fb5e8dbe1bc519caf75c26975002bf5c62291a199
cleanmarkers
mark-complete: 0
"""

def partition_script(partition_spec, partition_number):
    file, name, human_size, partition_eblocks = partition_spec
    print "writing ofw partition script for %s" % name
    script = "set-partition: %s\nmark-pending: 0\n" % name
    # for each eblock in the file, sha256 sum it
    file_eblocks = os.stat(file).st_size / eblock_bytes
    imgfile = open(file, "r")
    for block in range(0, file_eblocks):
        script += "eblock: %s sha256 %s\n" % (dec2hex(block),
                hashlib.sha256(imgfile.read(eblock_bytes)).hexdigest())
    script += "cleanmarkers\nmark-complete: 0\n"
    return script


from optparse import OptionParser

parser = OptionParser()
parser.add_option("-f", "--file", dest="nand_image_file", default="nand.img",
                  help="use FILE as output image file", metavar="FILE")
parser.add_option("-d", "--data-file", dest="data_file", default="data.img",
        help="write OFW update script to FILE.  default: data.img", metavar="FILE")
parser.add_option("-e", "--eblock-size", dest="eblock_size", 
                  default=str(128*KiB_bytes),
                  help="use SIZE bytes as eblock size for sha256 sums.  default 128KiB",
                  metavar="SIZE")
parser.add_option("-p", "--partition-layout", dest="partition_layout",
        help="set LAYOUT as partition layout.  Format is [<tree> <intermediate image file name> <partition name> <size specifier> <compressor>] for each partition: e.g. 'bootfs/ bootfs.img boot 20MiB zlib  rootfs/ rootfs.img root -1 lzo'", metavar="'LAYOUT'")

(options, args) = parser.parse_args()

eblock_bytes = human_size_to_bytes(options.eblock_size)
if options.partition_layout:
    partition_layout = options.partition_layout.split()
else:
    print "you must specify a partition layout"
    exit(1)

print "using eblock_size = %i" % eblock_bytes
print "using data_file OFW update script %s" % options.data_file
if options.nand_image_file:
    print "using output image_file %s" % options.nand_image_file


# get partition specifications

print partition_layout

if not len(partition_layout) % 3 == 0:
    raise Exception("malformed partition layout")


partition_layout_string = ""
partitions = []
for i in range(0, len(partition_layout)/3):
    j = i*3
    file = partition_layout[j]
    name = partition_layout[j+1]
    human_size = partition_layout[j+2]
    eblocks = '-1' if human_size == '-1' else str(dec2hex(human_size_to_bytes(human_size) / eblock_bytes))
    print "file: %s, name: %s, size: %s" % (file, name, human_size)
    partitions.append([file, name, human_size, eblocks])
    partition_layout_string += "%s %s  " % (name, eblocks)

print partition_layout_string
#print partitions


# data.img ofw script header
ofw_updater_script = """data:  %s
erase-all
partitions:  %s
""" % (options.nand_image_file, partition_layout_string)

partition_number = 0
for partition in partitions:
    file, name, human_size, partition_eblocks = partition
    partition_number += 1
    ofw_updater_script += partition_script(partition, partition_number)
    
#print ofw_updater_script

dataf = open(options.data_file, 'w')
dataf.write(ofw_updater_script)
dataf.close()

# now concatenate the image files into the output file
if options.nand_image_file:
    print "building NAND image file %s from intermediates" % \
            options.nand_image_file
    os.system("cat %s >%s" % (' '.join(map(lambda x: x[0], partitions)),
                              options.nand_image_file))

print "done"