#!/usr/bin/env python
"""
Read TI COFF file into a python object
"""


from __future__ import print_function

from builtins import bytes, object, range
from array import array
from collections import namedtuple
from optparse import OptionParser
from struct import unpack, calcsize


# See file COPYING in this source tree
# See file ticoff.asi_license.txt in this source tree
__copyright__ = "\n".join(
    [
        'Copyright 2011 "Eliot Blennerhassett" <eblennerhassett@audioscience.com>',
        # https://gist.github.com/eliotb/1073231
        "Copyright 2016, EPC Power Corp."
        # Added Python3 support and a few tweaks
    ]
)
__license__ = "GPLv2+"


def read_struct(file, format):
    """read struct data formatted according to format"""
    datalen = calcsize(format)
    data = file.read(datalen)
    if len(data) != datalen:
        raise EOFError
    ret = unpack(format, data)
    if len(ret) < 2:
        ret = ret[0]
    return ret


def read_cstr(file):
    """read zero terminated c string from file"""
    output = ""
    while True:
        char = file.read(1)
        if len(char) == 0:
            raise RuntimeError("EOF while reading cstr")
        char = chr(bytes(char)[0])
        if char == "\0":
            break
        output += char
    return output


# From spraao8.pdf table 7
class Section(
    namedtuple(
        "Section",
        "name, virt_size, virt_addr, raw_data_size, raw_data_ptr, relocation_ptr,"
        + "linenum_ptr, linenum_count, reloc_count, flags, reserved, mem_page, data",
    )
):

    __slots__ = ()  # prevent creation of instance dictionaries
    section_fmt = "<8s9L2H"
    section_flags = [
        (0x00000001, "STYP_DSECT"),
        (0x00000002, "STYP_NOLOAD"),
        (0x00000004, "STYP_GROUPED"),
        (0x00000008, "STYP_PAD"),
        (0x00000010, "STYP_COPY"),
        (0x00000020, "STYP_TEXT"),
        (0x00000040, "STYP_DATA"),
        (0x00000080, "STYP_BSS"),
        (0x00000100, "STYP_100"),
        (0x00000200, "STYP_200"),
        (0x00000200, "STYP_400"),
        (0x00000800, "STYP_800"),
        (0x00001000, "STYP_BLOCK"),
        (0x00002000, "STYP_PASS"),
        (0x00004000, "STYP_CLINK"),
        (0x00008000, "STYP_VECTOR"),
        (0x00010000, "STYP_PADDED"),
    ]

    @property
    def is_zero_filled(self):
        return (self.data is not None) and not any([d != "\0" for d in self.data])

    def fill_value(self):
        if self.data is None:
            return None

        if len(self.data) % 4:
            return None  # only handle 32 bit fill

        a = array("I", self.data)
        if all([d == a[0] for d in a]):
            return a[0]
        else:
            return None

    @property
    def is_loadable(self):
        return not (self.flags & 0x00000010 or self.data is None)

    def __str__(self):
        s = "Section %s" % self.name
        if self.flags:
            s += " Flags ="
            for f in Section.section_flags:
                if self.flags & f[0]:
                    s = s + " " + f[1]
        for fv in self._asdict().items():
            if fv[0] == "data" or fv[0] == "name":
                continue
            s = s + "\n\t%s = 0x%X" % fv
        return s


class Symbol(
    namedtuple(
        "Symbol",
        "name, value, section_number, reserved," "storage_class, number_of_aux_entries",
    )
):

    __slots__ = ()  # prevent creation of instance dictionaries
    symbol_fmt = "<8sLhHcc"
    symbol_auxiliary_format = "<LHH10c"
    symbol_flags = [
        (0, "C_NULL"),
        (1, "C_AUTO"),
        (2, "C_EXT"),
        (3, "C_STAT"),
        (4, "C_REG"),
        (5, "C_EXTREF"),
        (6, "C_LABEL"),
        (7, "C_ULABEL"),
        (8, "C_MOS"),
        (9, "C_ARG"),
        (10, "C_STRTAG"),
        (11, "C_MOU"),
        (12, "C_UNTAG"),
        (13, "C_TPDEF"),
        (14, "C_USTATIC"),
        (15, "C_ENTAG"),
        (16, "C_MOE"),
        (17, "C_REGPARM"),
        (18, "C_FIELD"),
        (19, "C_UEXT"),
        (20, "C_STATLAB"),
        (21, "C_EXTLAB"),
        (27, "C_VARARG"),
        (100, "C_BLOCK"),
        (101, "C_FCN"),
        (102, "C_EOS"),
        (103, "C_FILE"),
        (104, "C_LINE"),
    ]


class Coff(object):
    Header = namedtuple(
        "Header",
        "machine_type, section_count, time_date, symbol_table_ptr, "
        + "symbol_count, optional_header_size, flags, target_id",
    )
    header_fmt = "<2H3L3H"

    OptionalHeader = namedtuple(
        "OptionalHeader",
        "magic, version, exe_size, init_data_size, uninit_data_size, entry_point, exe_start, init_start",
    )
    optheader_fmt = "<2H6L"

    def __init__(self, filename=None):
        self.header = None
        self.optheader = None
        self.sections = []
        if filename is not None:
            self.from_file(filename)

    def from_file(self, name):
        with open(name, "rb") as f:
            self.from_stream(f)

    def from_stream(self, f):
        self.header = self.Header(*read_struct(f, self.header_fmt))
        self.optheader = self.OptionalHeader(*read_struct(f, self.optheader_fmt))
        self.sections = []
        for i in range(self.header.section_count):
            section = Section(*read_struct(f, Section.section_fmt), data=None)
            section = section._replace(name=self.symname(f, section.name))
            if section.raw_data_ptr and section.raw_data_size:
                here = f.tell()
                f.seek(section.raw_data_ptr)
                # TODO: `2 *` is hard coded to handle the 2-bytes per
                #       address scenario.  This should obviously be
                #       detected somehow, unless it is always correct.
                data = bytes(f.read(2 * section.raw_data_size))
                f.seek(here)
                section = section._replace(data=data)
            self.sections.append(section)

        self.symbols = []
        f.seek(self.header.symbol_table_ptr)
        for i in range(self.header.symbol_count):
            symbol = Symbol(*read_struct(f, Symbol.symbol_fmt))
            try:
                symbol = symbol._replace(name=self.symname(f, symbol.name))
            except UnicodeDecodeError:
                # TODO: not sure what to do with these
                pass
            symbol = symbol._replace(
                number_of_aux_entries=symbol.number_of_aux_entries[0]
            )
            try:
                symbol = symbol._replace(
                    storage_class=next(
                        f
                        for f in symbol.symbol_flags
                        if symbol.storage_class[0] == f[0]
                    )
                )
            except StopIteration:
                print("bad {}".format(symbol))
            self.symbols.append(symbol)

            # TODO: this breaks by hitting EOF but it seems we would need
            #       to read the aux entries somewhere
            # if symbol.number_of_aux_entries:
            #     read_struct(f, Symbol.symbol_auxiliary_format)

        stack_section_number = next(
            i for i, s in enumerate(self.sections) if s.name == ".stack"
        )
        self.variables = [
            str(s)
            for s in self.symbols
            if isinstance(s.name, str)
            and s.name.startswith("_")
            and s.section_number == stack_section_number
        ]

        self.entry_point = self.optheader.entry_point
        self.sections.sort(key=lambda s: s.virt_addr)

    @property
    def loadable_sections(self):
        return [s for s in self.sections if s.is_loadable]

    def string_table_entry(self, f, offset):
        here = f.tell()
        seek = self.header.symbol_table_ptr + self.header.symbol_count * 18 + offset
        f.seek(seek)
        s = read_cstr(f)
        f.seek(here)
        return s

    def symname(self, f, value):
        parts = unpack("<2L", value)
        if parts[0] == 0:
            return self.string_table_entry(f, parts[1])
        else:
            return str(value.decode("ascii").rstrip("\0"))

    def __str__(self):
        s = "TI coff file"
        if self.header is not None:
            for fv in self.header._asdict().items():
                s = s + "\n\t%s = 0x%X" % fv
        if self.optheader is not None:
            s += "\nOptional Header"
            for fv in self.optheader._asdict().items():
                s = s + "\n\t%s = 0x%X" % fv
        return s


if __name__ == "__main__":
    parser = OptionParser()
    parser.add_option(
        "-d",
        "--dump_data",
        action="store_true",
        dest="dump",
        default=False,
        help="dump section data",
    )

    options, args = parser.parse_args()

    c = Coff(args[0])
    print(c)

    for section in c.sections:
        print(section)
        if section.data is not None and options.dump:
            print(repr(section))
