
import mmap
import os
import sys
import struct
import imghdr

def get_image_size(fname):
    '''Determine the image type of fhandle and return its size.
    from draco'''
    with open(fname, 'rb') as fhandle:
        head = fhandle.read(24)
        if len(head) != 24:
            return
        if imghdr.what(fname) == 'png':
            check = struct.unpack('>i', head[4:8])[0]
            if check != 0x0d0a1a0a:
                return
            width, height = struct.unpack('>ii', head[16:24])
        elif imghdr.what(fname) == 'gif':
            width, height = struct.unpack('<HH', head[6:10])
        elif imghdr.what(fname) == 'jpeg':
            try:
                fhandle.seek(0) # Read 0xff next
                size = 2
                ftype = 0
                while not 0xc0 <= ftype <= 0xcf:
                    fhandle.seek(size, 1)
                    byte = fhandle.read(1)
                    while ord(byte) == 0xff:
                        byte = fhandle.read(1)
                    ftype = ord(byte)
                    size = struct.unpack('>H', fhandle.read(2))[0] - 2
                # We are at a SOFn block
                fhandle.seek(1, 1)  # Skip `precision' byte.
                height, width = struct.unpack('>HH', fhandle.read(4))
            except Exception: #IGNORE:W0703
                return
        else:
            return
        return width, height

if (len(sys.argv) < 3 and len(sys.argv) > 4 ) or not os.path.isfile(sys.argv[1]) or sys.argv[2] not in ("e", "r"):
    print("usage : %s img_file.bin e|r [-f]\n\te : extract\n\tr : rebuild\n\t-f : force" % sys.argv[0])
    os._exit(1)

if len(sys.argv) == 4 and not sys.argv[3] == "-f":
    print("usage : %s img_file.bin e|r [-f]\n\te : extract\n\tr : rebuild\n\t-f : force" % sys.argv[0])
    os._exit(1)

infile = sys.argv[1]
mode = sys.argv[2]
force = len(sys.argv) == 4

with open(infile, "r+b") as f:
    mm = mmap.mmap(f.fileno(), 0)
    
    num_img = struct.unpack("<L", mm[0:4])[0]
    print("num of images : %d" % num_img)

    basename = os.path.splitext(infile)[0]

if mode == "e":

    for i in range(0, num_img):
        print("extracting image : %d" % (i + 1))
        header_offset = 4 + i * 20
        size_x, size_y, unknow_1, unknow_2, size, file_offset = struct.unpack("<HHIIII", mm[header_offset: header_offset + 20])
        print("\timage size : %dx%d" % (size_x, size_y))
        print("\tunknow 1 : 0x%x" % (unknow_1))
        print("\tunknow 2 : %d" % (unknow_2))
        print("\tfile size : 0x%x / %d" % (size, size))
        print("\tfile offset : 0x%x / %d" % (file_offset, file_offset))
        fname = "%s_%d.png" % (basename, i + 1)
        if os.path.exists(fname) and not force:
            print("file exists use force (-f) to overwrite : " + fname)
            os._exit(1)
        with open(fname, "wb") as dumpf:
            dumpf.write(mm[file_offset : file_offset + size])
        print("\tdumped in %s" % fname)

else:
    outheader = b""
    outheader += struct.pack("<L", num_img)
    new_file_offset = 4 + (20 * num_img)
    outdata = b""
    for i in range(0, num_img):
        print("image : %d" % (i + 1))
        header_offset = 4 + i * 20
        size_x, size_y, unknow_1, unknow_2, size, file_offset = struct.unpack("<HHIIII", mm[header_offset: header_offset + 20])
        print("\timage size : %dx%d" % (size_x, size_y))
        print("\tunknow 1 : 0x%x" % (unknow_1))
        print("\tunknow 2 : %d" % (unknow_2))
        print("\tfile size : 0x%x / %d" % (size, size))
        print("\tfile offset : 0x%x / %d" % (file_offset, file_offset))
        fname = "%s_%d.png" % (basename, i + 1)
        if not os.path.exists(fname):
            print("Error : filename doest not exists, use extract first : " + fname)
            os._exit(1)

        w, h = get_image_size(fname)
        print("\tdisk image size : %dx%d" % (w, h))

        if not force and ( w != size_x or h != size_y ):
            print("image size on disk is different then in header, use -f to force")
            os._exit(1)

        disk_file_size = os.path.getsize(fname)
        print("\tdisk file size : 0x%x / %d" % (disk_file_size, disk_file_size))
        print("\tnew file offset : 0x%x / %d" % (new_file_offset, new_file_offset))
        outheader += struct.pack("<HHIIII", w, h, unknow_1, unknow_2, disk_file_size, new_file_offset)
        new_file_offset += disk_file_size
        outdata += open(fname, "rb").read()

    print("size header : %d" % len(outheader))
    print("size data : %d" % len(outdata))
    new_file = basename + "_new.bin"
    open(new_file, "wb").write(outheader + outdata)
    print("New file created in : " + new_file)
    
mm.close()
