import sys
import os

if len(sys.argv) != 2:
    print(f'usage: {sys.argv[0]} <dumpfile with oob and any software crc removed> <outputfilename,optional>')
    exit(1)
fname = sys.argv[1]

if len(sys.argv) == 3:
    fixed = sys.argv[2]
else:
    fixed = "output.bin"

verbose = True
pageSize = 0x840
pagesPerBlock = 64
blockSize = pageSize*pagesPerBlock

def ubiParse():
    with open(fname, 'rb') as fin:
        with open(fixed, 'wb') as fout:
            #fin.read(0xBE00000)
            while True:
                ec = fin.read(pageSize)
                vid = fin.read(pageSize)
                if len(ec) == 0 or len(vid) == 0:
                    break
                #print(ec[0:4],vid[0:4])
                
                if b"\x55\x42\x49\x23" == ec[0:4] and b"\x55\x42\x49\x21" == vid[0:4]:
                    #print(ec[0:4],vid[0:4])
                    if verbose:
                        print("Found EC and VID headers")
                    dataSize = int.from_bytes(vid[20:24], byteorder='big')
                    if verbose:
                        print("Datasize", dataSize)
                    if dataSize != 0 and dataSize < 0xFFFFFF:
                        #continue
                        fout.write(ec)
                        fout.write(vid)
                        fout.write((fin.read(dataSize)))
                        if (dataSize/0x800)%2 == 1:
                            fin.read(0x800)
                        
                        print("Wrote", dataSize, "bytes to file")
                #break
    print("Done")


def EFSParse():
    with open(fname, 'rb') as fin:
        with open(fixed, 'wb') as fout:
            while True:
                page = fin.read(pageSize)
                if len(page) == 0:
                    break
                    
                if b"EFSSuper" == page:
                    if verbose:
                        print("EFS Superblock")
                        print("Block Size:", int.from_bytes(vid[10:13]))
                        print("Block Size:", int.from_bytes(vid[10:13]))
                    dataSize = int.from_bytes(vid[20:24], byteorder='big')
                    if verbose:
                        print("Datasize", dataSize)
                    if dataSize != 0 and dataSize < 0xFFFFFF:
                        #continue
                        fout.write(ec)
                        fout.write(vid)
                        fout.write((fin.read(dataSize)))
                        if (dataSize/0x800)%2 == 1:
                            fin.read(0x800)
                        
                        #print("Wrote", dataSize, "bytes to file")
                #break
    print("Done")


def MIBParse():
    partitions = []
    extract = False
    with open(fname, 'rb') as fin:
        while True:
            page = fin.read(pageSize)
            
            if len(page) == 0:
                break
                
            if b"\xaa\x73\xee\x55" == page[0:4]:
                numEntries = int.from_bytes(page[12:16], byteorder="little")
                if verbose:
                    print("MIB Partition Table")
                    print("Number of Entries:", numEntries)
                for i in range(16,16+numEntries*26,26):
                    name = page[i:i+16].decode("utf8")
                    offset = int.from_bytes(page[i+16:i+16+4], byteorder="little")
                    length = int.from_bytes(page[i+16+4:i+16+4+4], byteorder="little")
                    partitions.append([name,offset*blockSize,length*blockSize])
                    if verbose:
                        print("Partition Name:", name)
                        print("Partititon Offset from 0:", hex(offset*blockSize))
                        print("Partition Length:", hex(length*blockSize))
                        print("Partition ECC config 0=Main, FF=Main+Spare:", page[i+25:i+26])
                        print("")
        if extract:            
            print("Partition table found, extracting the partitions to separate files")
            if not os.path.exists("UE910-extracted"):
                os.makedirs("UE910-extracted")
            for part in partitions:
                fin.seek(part[1])
                with open("UE910-extracted/"+part[0][2:].split("\x00")[0], "wb") as out:
                    out.write(fin.read(part[2]))
                    print("File",part[0][2:],"written to folder extracted")
            
    
MIBParse()