#!/usr/bin/python
import sys
import os
import tempfile
import getopt
from struct import *
from ctypes import *
from idicore.idicore import *

verbose = False

def parse_varlen_number(ls):
    l = 0
    while len(ls) > 0:
        c = ls[0]
        ls = ls[1:]
        pack = unpack("B", c)
        b = pack[0]
        l = (l << 8) + b
    return l


class DebugFrameDecopression:

    def __init__(self):
        self.cie_num_to_offset = {}
        self.fde_num_to_data = {}
        self.next_cie_num = 1
        self.next_fde_num = 1

    def decompress_cie(self, b):
        if verbose:
            print "Decompressing a CIE number %d" % self.next_cie_num
        assert (b & 0x40) != 0
        
        sl = b & 0x3f
        if (sl == 0):
            lenght = read_uleb128(self.inf)
        else:
            length = sl
            
        data = self.inf.read(length)

        self.cie_num_to_offset[self.next_cie_num] = self.outf.tell()
        write_word(self.outf, length + address_size)
        write_word(self.outf, cie_id)
        self.outf.write(data)
        self.next_cie_num = self.next_cie_num + 1

        if verbose:
            print "    length: 0x%x" % (length + address_size)

    def decompress_new_fde(self, hdr_b):
        if verbose:
            print "Decompressing a new FDE number %d" % self.next_fde_num
        sc = (hdr_b & 0x3f) >> 3
        sl = (hdr_b & 0x7) + 1
        
        if (sc == 0):
            cie = read_uleb128(self.inf)
        else:
            cie = sc

        ls = self.inf.read(sl)
        length = parse_varlen_number (ls)
        addr = read_address(self.inf)
        data = self.inf.read(length)
        
        self.fde_num_to_data[self.next_fde_num] = data

        assert self.cie_num_to_offset.has_key(cie)
        write_word(self.outf, length + address_size + word_size)
        write_word(self.outf, self.cie_num_to_offset[cie])
        write_address (self.outf, addr)
        self.outf.write (data)
        self.next_fde_num = self.next_fde_num + 1

        if verbose:
            print ("    lenght: 0x%x, cie_offset: 0x%x, address: 0x%x"
                   % ((length + address_size + word_size), 
                      self.cie_num_to_offset[cie], addr))

    def decompress_fde_ref(self, hdr_b):
        if verbose:
            print "Decompressing an FDE reference"
        sc = (hdr_b & 0x3f) >> 3
        sl = (hdr_b & 0x7) + 1

        if (sc == 0):
            cie = read_uleb128(self.inf)
        else:
            cie = sc
        
        lenstr = self.inf.read(sl)
        fde_num = parse_varlen_number(lenstr)
        addr = read_address(self.inf)
        data = self.fde_num_to_data[fde_num]

        write_word(self.outf, len(data) + address_size + word_size)
        write_word(self.outf, self.cie_num_to_offset[cie])
        write_address (self.outf, addr)
        self.outf.write (data)

        if verbose:
            print ("    ref: %d, lenght: 0x%x, cie_offset: 0x%x, address: 0x%x"
                   % (fde_num, (len(data) + address_size + word_size), 
                      self.cie_num_to_offset[cie], addr))
        
                

    def process_file(self, inf, outf):
        inf_pos = 0
        outf_pos = 0
        try:
            try:
                self.inf = inf
                self.outf = outf
            
                while True:
                    inf_pos = self.inf.tell()
                    outf_pos = self.outf.tell()

                    if verbose:
                        print ("\ninput offset: 0x%x, output offset: 0x%x"
                               % (inf_pos, outf_pos))

                    try:
                        b = read_byte(self.inf)
                    except EEOF:
                        break                

                    if (b & 0x80) != 0:
                        self.decompress_cie (b)
                    else:
                        if (b & 0x40) != 0:
                            self.decompress_new_fde(b)
                        else:
                            self.decompress_fde_ref(b)
            except:
                print "========================================"
                print "Unhandled exception raised at offsets:"
                print ("inf: %d (0x%x), outf %d (0x%x)" 
                       % (inf_pos, inf_pos, outf_pos, outf_pos))
                print "========================================"
                raise
        finally:
            self.inf = None
            self.outf = None



def skip_cies_at_the_end(ref):
    if verbose:
        print ("Skipping CIEs at the end of the reference file (pos: 0x%x)"
               % ref.tell())
    while True:
        try:
            elen = read_word(ref)
        except EEOF:
            return
        
        id = read_word(inf)
        elen = elen - word_size
        
        assert id == cie_id, "Decompressed file ended prematurely"
        ref.seek(elen, 1)
        if verbose:
            print "    skipping 0x%x bytes" % (elen + 2 * word_size)


def check_skip_cie(ref_data, ref_pos, known_cies, cie_num_to_offset, 
                   cie_offset_to_offset):
    id = -1
    i = 0
    for past in known_cies:
        if past == ref_data:
            id = i
            break
        i = i + 1
        
    assert id >= 0, "Stumbled upon an unknown CIE in the reference file"
    past_offset = cie_num_to_offset[id]
    cie_offset_to_offset[ref_pos] = past_offset
    
    if verbose:
        print "Skipped a CIE that is the same as CIE number %d" % id


def compare_information(tst, ref):
    cie_offset_to_offset = {}
    cie_num_to_offset = {}
    known_cies = []
    cies_num = 0
    tst_pos = 0
    ref_pos = 0
    tst_entry_loaded = False

    try:
        while True:
            if (not tst_entry_loaded):
                tst_pos = tst.tell()
                try:
                    tst_elen = read_word(tst)
                except EEOF:
                    skip_cies_at_the_end(ref)
                    return
                tst_id = read_word(tst)
                tst_elen = tst_elen - word_size
                tst_data = tst.read(tst_elen)

            ref_pos = ref.tell()
            try:
                ref_elen = read_word(ref)
            except:
                assert False, "Decompressed file too long"
            ref_id = read_word(ref)
            ref_elen = ref_elen - word_size
            ref_data = ref.read(ref_elen)

            if verbose:
                print ("\ntest offset: 0x%x, reference offset: 0x%x"
                       % (tst_pos, ref_pos))

            if (tst_id == cie_id and ref_id == cie_id):
                if (tst_data == ref_data):
                    known_cies.append(ref_data)
                    cie_num_to_offset[cies_num] = tst_pos
                    cie_offset_to_offset[ref_pos] = tst_pos
                    if verbose:
                        print "New CIE number %d found" % cies_num
                    cies_num = cies_num + 1
                    tst_entry_loaded = False
                else:
                    check_skip_cie(ref_data, ref_pos, known_cies, 
                                   cie_num_to_offset, cie_offset_to_offset)
                    tst_entry_loaded = True

            elif (tst_id != cie_id and ref_id == cie_id):
                check_skip_cie(ref_data, ref_pos, known_cies, 
                               cie_num_to_offset, cie_offset_to_offset)
                tst_entry_loaded = True
            elif (tst_id != cie_id and ref_id != cie_id):
                assert tst_id == cie_offset_to_offset[ref_id], "CIE mismatch"
                assert tst_data == ref_data, "FDE data mismatch"

                if verbose:
                    print "Found matching FDEs"

                tst_entry_loaded = False
            else:    
                # tst_id == cie_id and ref_id != cie_id
                assert False, "Unpaired CIE in the tst file"

    except:
        print "========================================"
        print "Unhandled exception raised at offsets:"
        print ("tst: %d (0x%x), ref %d (0x%x)" 
               % (tst_pos, tst_pos, ref_pos, ref_pos))
        print "========================================"
        raise
        
def print_usage():
    print """    deframe.py [-rxtvh] [-o output] files... 
        - uncompress or test compressed unwind debug info

This utility either decompresses (by default) or tests (when given
switch -t or --test) unwind debug info compressed by frame.py.  When
decompressing, the raw decompressed output is stored into an output
file specified by the -o option (the default is
uncompressed_unwind_info, potentially with a suffix).  When testing,
data are decompressed to a temporary file which is then intelligently
compared with the original .debug_frame info in the elf file.

-r (or --raw) The input file is not an elf file but raw binary
   compresssed data.

-x (or --64) The data are relative to a 64bit elf file (only
   meaningful when used with the -r option)

-t (or --test) Switch to the test mode as described above.

-v produces a lot of verbose debugging output.

-o filename (or --output=filename) Stores the decompressed output in
   the given file.  This option has no effect in the test mode.

-h (or --help) displays this message


"""
            
def main():
    """The main decompression and testing script driver function."""

    global verbose 

    if (len(sys.argv) < 2):
        die("You need to specify at least one input file.")

    try:
        opts, args = getopt.gnu_getopt(sys.argv[1:], "xro:tv", 
                                       ["64", "raw", "output=", "test",
                                        "lists"])
    except getopt.GetoptError, err:
        # print help information and exit:
        print_usage()
        die(str(err)) # will print something like "option -a not recognized"
    outf_filename_base = "uncompressed_unwind_info"
    test = False
    raw = False
    verbose = False
    keep_temp_files = False
    switch_to_32()
    for o, a in opts:
        if o in ( "-o", "--output"):
            outf_filename_base = a
        elif o in ("-r", "--raw"):
            raw = True
        elif o in ("-t", "--test"):
            test = True
        elif o in ("-x", "--64"):
            switch_to_64()
        elif o  == "-v":
            verbose = True
        elif o in ("-h", "--help"):
            print_usage()
            sys.exit(1)
        else:
            print_usage()
            assert False, "Unhandled option"

    filenum = 0
    for argument in args:
        if verbose:
            print "------ Processing argument:", argument, " ------"

        if not raw:
            identify_elf_file(argument)
            inf_filename = tempfile.mktemp("frame-comp", temp_prefix)
            cmd_str_extr_our = ("objcopy -O binary --set-section-flags " + 
                                "%s=alloc -j %s %s %s" 
                                % (compressed_unwind_section_name,
                                   compressed_unwind_section_name,
                                   argument, inf_filename))
            if verbose:
                print "| Running", cmd_str_extr_our
            if os.system(cmd_str_extr_our) != 0:
                die(("Could not extract %s section from file %s\n")
                    % (compressed_unwind_section_name, argument))
        else:
            inf_filename = argument

        if test:
            ref_filename = tempfile.mktemp("frame-ref", temp_prefix)
            cmd_str_extr_orig = ("objcopy -O binary --set-section-flags " + 
                                 ".debug_frame=alloc -j .debug_frame %s %s" 
                                 % (argument, ref_filename))
            if verbose:
                print "| Running", cmd_str_extr_orig
            if os.system(cmd_str_extr_orig) != 0:
                die("Could not extract .debug_frame from file %s\n" % elfname)

        inf = open(inf_filename, 'r')
        try:
            if test:
                outf_fd, outf_filename = tempfile.mkstemp("frame-decomp", 
                                                        temp_prefix)
                outf = os.fdopen(outf_fd, "r+")
            else:
                if filenum > 0:
                    outf_filename = "%s-%d" % (outf_filename_base, filenum)
                    filenum = filenum + 1
                else:
                    outf_filename = outf_filename_base
                    filenum = 2
                    outf = open(outf_filename, "w")
        except:
            inf.close()
            if not raw:
                os.remove(inf_filename)
            if test:
                os.remove(ref_filename)
            raise
        
        try:
            if verbose:
                print "* Decompression %s -> %s" % (inf_filename, outf_filename)
            frame = DebugFrameDecopression()
            frame.process_file(inf, outf)

            if test:
                if verbose:
                    print "* Comparing decompressed and reference info"
                outf.seek(0)
                ref = open(ref_filename, "r")
                try:
                    compare_information(outf, ref)

                finally:
                    ref.close()
        finally:
            inf.close()
            if not raw:
                os.remove(inf_filename)
            outf.close()
            if test:
                os.remove(outf_filename)
                os.remove(ref_filename)




if __name__ == '__main__':
    main()
