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

verbose = False
keep_temp_files = False

def add_to_list_dict(data, dict, max):
    """Finf or add the data with hash h to a dictionary of lists.

    D is string of actual data
    DICT is the dictionary of lists of (data, num) tuples.
    MAX is the number that should be assigned to new data.

    If the data is not present in the dictionary, an appropriate tuple
    in the right list is created for it and fucntion returns zero.
    Otherwise, the number stored in the tuple is created.
    """
    assert max > 0
    hash = jenkins_hash(data)
    if (dict.has_key (hash)):
        list = dict[hash]
        for t in list:
            if (t[0] == data):
                assert t[1] > 0
                return t[1]
        list.append((data, max))
        return 0
    else:
        dict[hash] = [(data, max)];
        return 0

def gen_varlen_number(number):
    s = number
    data = pack("B", (number & 0xff))
    z = number >> 8;
    sl = 0;
    curmax = 256;
    while s >= curmax:
        assert z > 0
        curmax = curmax << 8
        sl = sl + 1

        data = pack("B", z & 0xff) + data
        z = z >> 8
    assert sl < 8
    assert z == 0
    return (sl, data)

  
class DebugFrameProcessing:

    def __init__(self):
        self.cie_next_num = 1
        self.fde_next_num = 1
        self.cie_dict = {}
        self.fde_dict = {}
        self.cie_num_to_offset = {}
        self.cie_offset_to_num = {}

    def write_compressed_cie(self, data, size):
        """Write a compressed CIE in DATA with size SIZE"""

        if verbose:
            print ("    writing a CIE %d to output" % self.cie_next_num)
        
        if (size > 63):
            sl = 0
        else:
            sl = size

        b = 0xC0 | sl
        write_byte (self.outf, b)
        
        if (sl == 0):
            write_uleb_128(self.outf, size)

        self.outf.write(data)
        return
        

    def process_cie(self, pos, size):
        """Process input CIE in inf at position POS which is SIZE long
        
        The current inf pointer must be right after the CIE ID
        SIZE id the length of the CIE excluding its length as it is
        stored in the file and the CIE ID
        """
        
        if verbose:
            print "Processing a CIE"
        data = self.inf.read(size)
        num = add_to_list_dict(data, self.cie_dict, self.cie_next_num);
        if (num == 0):
            self.write_compressed_cie (data, size)
            self.cie_num_to_offset[self.cie_next_num] = self.outf.tell()
            self.cie_offset_to_num[pos] = self.cie_next_num
            self.cie_next_num = self.cie_next_num + 1
        else:
            self.cie_offset_to_num[pos] = num
                
            
    def write_new_fde(self, cid, addr, data, size):
        if verbose:
            print ("    writing a new FDE %d (for CIE %d, addr 0x%x)"
                   % (self.fde_next_num, cid, addr))
        if (cid > 7):
            scid = 0
        else:
            scid = cid

        len = gen_varlen_number(size)
        
        v = 0x40 | (scid << 3) | len[0]
        write_byte(self.outf, v)
        if (scid == 0):
            write_uleb_128(self.outf, cid)
        
        self.outf.write(len[1])
        write_address(self.outf, addr)
        self.outf.write(data)

    def write_fde_ref(self, cid, num, addr):
        if verbose:
            print ("    writing an FDE reference to %d (for CIE %d, addr 0x%x)"
                   % (num, cid, addr))

        if (cid > 7):
            scid = 0
        else:
            scid = cid
        
        num_len, num_val = gen_varlen_number(num)
        #print "number (length and dump):", num_len
        #print debug_hexdump(num_val)#!!!
        v = (scid << 3) | num_len
        assert v < 0x40
        write_byte(self.outf, v)
        if (scid == 0):
            write_uleb_128(self.outf, cid)

        self.outf.write(num_val)
        write_address(self.outf, addr)

    
    def process_fde(self, cie, pos, size):
        if verbose:
            print "Processing an FDE"
        addr = read_word(self.inf)
        size = size - word_size

        #print "------------------------------"

        assert (self.cie_offset_to_num.has_key(cie))
        cid = self.cie_offset_to_num[cie]

        data = self.inf.read(size);
        num = add_to_list_dict(data, self.fde_dict, self.fde_next_num);
        if (num == 0):
            self.write_new_fde(cid, addr, data, size)
            self.fde_next_num = self.fde_next_num + 1
        else:
            self.write_fde_ref(cid, num, addr)
            

    def process_raw_section_file(self, inf, outf):
        """Compress the raw unwind info in file INF and store it to OUTF

        This function is basically a loop that reads all entries and
        calls appropriate methods on them.
        """

        self.inf = inf
        self.outf = outf
        pos = 0
        try:
            while True:
                pos = inf.tell()
                if verbose:
                    print ("\ninput offset: 0x%x, output offset: 0x%x"
                           % (pos, outf.tell()))
                try:
                    entry_length = read_word (inf)
                except EEOF:
                    break
                load_size = entry_length

                id = read_word(inf)
                load_size = load_size - word_size

                if id == cie_id:
                    self.process_cie(pos, load_size)
                else:
                    self.process_fde(id, pos, load_size)
                
        finally:
            self.inf = None
            self.outf = None
        return pos


def process_elf(elfname):
    """Process the unwind information of an elf file."""
    global frame

    identify_elf_file(elfname)
    infilename = tempfile.mktemp("frame-in", temp_prefix)
    cmd_str_extract = ("objcopy -O binary --set-section-flags " + 
                       ".debug_frame=alloc -j .debug_frame %s %s" 
                       % (elfname, infilename))
    if verbose:
        print "| Running", cmd_str_extract
    if os.system(cmd_str_extract) != 0:
        sys.stderr.write("Could not extract .debug_frame section from " +
                         "file %s, skipping...\n" % elfname)
        return

    inf = open(infilename, 'r')
    try:
        try: 
            outf_fd, outf_name = tempfile.mkstemp("frame-out", temp_prefix)
            outf = os.fdopen(outf_fd, "w")
        except:
            inf.close()
            raise

        try:
            if verbose:
                print "* Compression %s -> %s" % (infilename, outf_name)
            frame = DebugFrameProcessing()
            processed = frame.process_raw_section_file(inf, outf)
        finally:
            inf.close()
            outf.close()

        if processed > 0:
            cmd_str_insert = ("objcopy --add-section %s=%s %s" 
                              % (compressed_unwind_section_name,
                                 outf_name, elfname))
            if verbose:
                print "| Running", cmd_str_insert
            if os.system(cmd_str_insert) != 0:
                sys.stderr.write("Could not add section %s to file %s\n"
                                 % (compressed_unwind_section_name, elfname))
                return
        
            
    finally:
        os.remove(infilename)
        if not keep_temp_files:
            os.remove(outf_name)
        else:
            print ("NOT removing temporary file %s" % outf_name)

def process_files_in_list_file(list):
    """Driving loop that takes in a list of files in a file."""

    if (len(sys.argv) < 2):
        die("You need to specify an input file.")
    list_file = open(list)
    try:
        for argument in list_file:
            argument = argument[:-1]
            if verbose:
                print "---- Processing elf file:", argument, " ----"
            process_elf(argument)
    finally:
        list_file.close()


def print_usage():
    print """    frame.py [-lvkh] arguments - compress debug unwind info

This utility compresses the .debug_frame section of elf files and
stores them back as %s sections.  The arguments are individual elf
files unless the -l (or --lists) switch is used, in that case they are
files containing lists of elf files, one per line.

-k (or --keep-temp) makes this tool keep the temporary file with the raw
   compressed data instead of deleting them.

-v produces a lot of verbose debugging output.

-h (or --help) displays this message

""" % compressed_unwind_section_name

def main():
    "The mainfunction parsing command line and driving appropriate compression"
    global verbose, keep_temp_files
    if (len(sys.argv) < 2):
        die("You need to specify at least one input file.")

    try:
        opts, args = getopt.gnu_getopt(sys.argv[1:], "lvkh", 
                                       ["lists", "keep-temp", "help"])
    except getopt.GetoptError, err:
        # print help information and exit:
        print_usage()
        die(str(err))
    verbose = False
    lists = False
    keep_temp_files = False
    for o, a in opts:
        if o in ( "-l", "--lists"):
            lists = True
        elif o in ("-k", "--keep-temp"):
            keep_temp_files = True
        elif o  == "-v":
            verbose = True
        elif o in ("-h", "--help"):
            print_usage()
            sys.exit(1)
        else:
            print_usage()
            assert False, "Unhandled option"

    for argument in args:
        if verbose:
            print "------ Processing argument:", argument, " ------"
        if lists:
            process_files_in_list_file(argument)
        else:
            process_elf(argument)






if __name__ == '__main__':
    main()
