# an example of a file loader in Python
# The scripting loader must define at least two functions: accept_file and load_file
# other optional functions are: save_file, move_segm, ...
#
# see also loader.hpp

import idaapi
from idc import *
import struct

ROM_SIGNATURE_OFFSET = 64
ROM_SIGNATURE        = "ECEC"
RomFormatName        = "Windows CE ROM"

# -----------------------------------------------------------------------
def dwordAt(li, off):
    li.seek(off)
    s = li.read(4)
    if len(s) < 4: 
        return 0
    return struct.unpack('<I', s)[0]

# -----------------------------------------------------------------------
def accept_file(li, n):
    """
    Check if the file is of supported format

    @param li: a file-like object which can be used to access the input data
    @param n : format number. The function will be called with incrementing 
               number until it returns zero
    @return: 0 - no more supported formats
             string "name" - format name to display in the chooser dialog
             dictionary { 'format': "name", 'options': integer }
               options: should be 1, possibly ORed with ACCEPT_FIRST (0x8000)
               to indicate preferred format
    """

    # we support only one format per file
    if n > 0:
        return 0

    # check the CECE signature
    li.seek(ROM_SIGNATURE_OFFSET)
    if li.read(4) == ROM_SIGNATURE:
        # accept the file
        return RomFormatName
    
    # unrecognized format
    return 0

# -----------------------------------------------------------------------
def load_file(li, neflags, format):
    
    """
    Load the file into database

    @param li: a file-like object which can be used to access the input data
    @param neflags: options selected by the user, see loader.hpp
    @return: 0-failure, 1-ok
    """
    
    if format == RomFormatName:
        jump = dwordAt(li, 0)
        if jump & 0xFF000000 == 0xEA000000: # looks like an ARM branch?
            idaapi.set_processor_type("arm", SETPROC_ALL|SETPROC_FATAL)
        
        li.seek(0, idaapi.SEEK_END)
        size = li.tell()

        #next dword after signature is a pointer to ROMHDR
        romhdr  = dwordAt(li, ROM_SIGNATURE_OFFSET + 4)
        
        # let's try to find such imagebase that potential ROMHDR's "physfirst" value matches it
        imgbase = (romhdr-size) & ~0xfff
        bases = []
        maxbase = 0
        while imgbase < romhdr+8:
            physfirst = dwordAt(li, romhdr - imgbase + 8)
            if physfirst == imgbase:
                bases.append(imgbase)
            imgbase += 0x1000

        if len(bases) == 1:
            start = bases[0]
        elif len(bases) > 1:
            print "Warning: several potential imagebases detemined: " + ", ".join("%08X"%i for i in bases)
            start = bases[-1]
        else:
            Warning("Unable to determine load image base.")
            start = 0x80000000
        print "Using imagebase %08X" % start
        
        physlast = dwordAt(li, romhdr - start + 12)
        if physlast <= start:
            Warning("Internal error")
            return 0
        size = physlast - start

        AddSeg(start, start+size, 0, 1, idaapi.saRelPara, idaapi.scPub)

        # copy bytes to the database
        li.seek(0)
        li.file2base(0, start, start+size, 0)

        idaapi.add_entry(start, start, "start", 1)
        print "Load OK"
        return 1

    Warning("Unknown format name: '%s'" % format)
    return 0

# -----------------------------------------------------------------------
def move_segm(frm, to, sz, fileformatname):
    Warning("move_segm(from=%s, to=%s, sz=%d, formatname=%s" % (hex(frm), hex(to), sz, fileformatname))
    return 0
