#!/usr/bin/env python

# nokia lzo image decompressor by ali1234
# usage: unlzo file.lzo

# how the format works
# rootfs.lzo is a sequence of compressed chunks
# each chunk begins with a 20 byte header which is made up of 5 unsigned 
# integers:

# magic_0: 0xb8c3b410
# magic_1: 0
# is_compressed: 1 if chunk is compressed, 0 if not
# length: length of chunk data in bytes (not including header)
# uncompressed_length: length of chunk after decompression

# so we read each chunk, check if it is compressed, and if so, pass it to lzo
# to decompress. (this involves putting some extra header onto the string)

# after the chunk is decompressed we have an uncompressed stream, but it still 
# isn't the original image. it contains only the used 512 byte blocks from 
# within the original image. There is a 0x2000 byte header containing what
# appears to be a checksum at 0x0 and the original filename of the image at
# 0x8-0x40 (null padded c string). After this comes a list of offsets and counts
# {uint, uint} which indicate where the used blocks go in the output file. This 
# goes up to 0x2000 and is null padded. After this comes 512 byte blocks of 
# data.

# To rebuild the original file, loop through the table of offsets and counts:
#    write 512 byte blocks of 0s to output until reaching offset
#    copy 512 byte blocks from the stream starting at 0x2000 until reaching
#        offset + count


import struct
import sys
import lzo


def readchunk(file):
    header = file.read(20)
    (a, b, c, d, e) = struct.unpack('IIIII', header)
    if a != 0xb8c3b410 or b != 0:
        print "Unexpected thing happened."
        exit(0)
    return (c,d,e,file.read(d))


class Unsparse(object):
    def __init__(self):
        self.current_block = 0
        self.blocklist = []
        self.blocklist_pos = 0
        self.setup = False

    def do_setup(self, data):
        print "First compressed chunk. Reading header."
        image_name, = struct.unpack('56s', data[0x8:0x40])
        image_name = image_name.strip('\x00')
        print image_name
        self.output = file(image_name, 'wb')
        for i in range(0x40,0x2000,0x8):
            a,b = struct.unpack('II', data[i:i+8])
            if a == b == 0:
                pass
            else:
                print a, b, "(", a+b, ")", hex((a+b)*512)
                self.blocklist.append((a,b))

    def handle_data(self, data):
        chunk_pos = 0
        if not self.setup:
            self.do_setup(data)
            self.setup = True
            chunk_pos = 0x2000

        while self.blocklist_pos < len(self.blocklist):
            offset,count = self.blocklist[self.blocklist_pos]
            while self.current_block < offset:
                self.output.write('\x00'*512)
                self.current_block += 1
            while self.current_block < (offset+count):
                self.output.write(data[chunk_pos:chunk_pos+512])
                self.current_block += 1
                chunk_pos += 512
                if chunk_pos >= len(data):
                    return
            self.blocklist_pos += 1
            print offset, count

     

infile = file(sys.argv[1], 'rb')
unsparse = Unsparse()
n = 0
while 1:
    c,d,e,data = readchunk(infile)
    if c == d == e == 0:
        print "Chunk", n, "is empty block. End of data."
        exit(0)
    print "Chunk", n, "\tat", hex(infile.tell()), "\tis compressed." if c else "not compressed.", "\tData size:", hex(d), "\tUncompressed size:", hex(e)
    if c:
        decomp = lzo.decompress(struct.pack(">BI", 0xf0, e)+data)
    else:
        decomp = data

    unsparse.handle_data(decomp)
    n += 1