#!/usr/bin/env python3
#
# Copyright (c) 2018, the Dart project authors.  Please see the AUTHORS file
# for details. All rights reserved. Use of this source code is governed by a
# BSD-style license that can be found in the LICENSE file.

import argparse
from ctypes import create_string_buffer
from struct import *

# FILE HEADER FLAGS
FILE_HEADER_RELFLG = 0x1  # No relocation information
FILE_HEADER_EXEC = 0x2  # Executable
FILE_HEADER_LNNO = 0x4  # No line number information
FILE_HEADER_LSYMS = 0x8  # Local symbols removed / not present
FILE_HEADER_AR32WR = 0x100  # File is 32-bit little endian

# SECTION HEADER FLAGS
SECTION_HEADER_TEXT = 0x20  # Contains executable code
SECTION_HEADER_DATA = 0x40  # Contains only initialized data
SECTION_HEADER_BSS = 0x80  # Contains uninitialized data
SECTION_HEADER_ALIGN32 = 0x00600000  # Align data on a 32-byte boundary
SECTION_HEADER_EXECUTE = 0x20000000  # The section can be executed as code
SECTION_HEADER_READ = 0x40000000  # The section can be read
SECTION_HEADER_WRITE = 0x80000000  # The section can be written to

# FILE HEADER FORMAT
# typedef struct {
#   unsigned short f_magic;         /* magic number             */
#   unsigned short f_nscns;         /* number of sections       */
#   unsigned long  f_timdat;        /* time & date stamp        */
#   unsigned long  f_symptr;        /* file pointer to symtab   */
#   unsigned long  f_nsyms;         /* number of symtab entries */
#   unsigned short f_opthdr;        /* sizeof(optional hdr)     */
#   unsigned short f_flags;         /* flags                    */
# } FILHDR;
FILE_HEADER_FORMAT = 'HHIIIHH'
FILE_HEADER_SIZE = calcsize(FILE_HEADER_FORMAT)
FILE_HEADER_MAGICS = {
    'x86': 0x014c,
    'x64': 0x8664,
    'arm': 0x1c0,
    'arm64': 0xaa64,
    'riscv32': 0x5032,
    'riscv64': 0x5064,
}
FILE_HEADER_NUM_SECTIONS = 1
FILE_HEADER_TIMESTAMP = 0
FILE_HEADER_SIZE_OF_OPTIONAL = 0
FILE_HEADER_FLAGS = FILE_HEADER_LNNO

# SECTION HEADER FORMAT
# typedef struct {
#   char           s_name[8];  /* section name                     */
#   unsigned long  s_paddr;    /* physical address, aliased s_nlib */
#   unsigned long  s_vaddr;    /* virtual address                  */
#   unsigned long  s_size;     /* section size                     */
#   unsigned long  s_scnptr;   /* file ptr to raw data for section */
#   unsigned long  s_relptr;   /* file ptr to relocation           */
#   unsigned long  s_lnnoptr;  /* file ptr to line numbers         */
#   unsigned short s_nreloc;   /* number of relocation entries     */
#   unsigned short s_nlnno;    /* number of line number entries    */
#   unsigned long  s_flags;    /* flags                            */
# } SCNHDR;
SECTION_HEADER_FORMAT = '8sIIIIIIHHI'
SECTION_HEADER_SIZE = calcsize(SECTION_HEADER_FORMAT)
SECTION_NAME_RDATA = b'.rdata'
SECTION_NAME_TEXT = b'.text'
SECTION_PADDR = 0x0
SECTION_VADDR = 0x0
SECTION_RAW_DATA_PTR = (
    FILE_HEADER_SIZE + FILE_HEADER_NUM_SECTIONS * SECTION_HEADER_SIZE)
SECTION_RELOCATION_PTR = 0x0
SECTION_LINE_NUMS_PTR = 0x0
SECTION_NUM_RELOCATION = 0
SECTION_NUM_LINE_NUMS = 0

# SYMBOL TABLE FORMAT
# typedef struct {
#   union {
#     char e_name[8];
#     struct {
#       unsigned long e_zeroes;
#       unsigned long e_offset;
#     } e;
#   } e;
#   unsigned long e_value;
#   short e_scnum;
#   unsigned short e_type;
#   unsigned char e_sclass;
#   unsigned char e_numaux;
# } SYMENT;
SYMBOL_TABLE_ENTRY_SHORT_LEN = 8
SYMBOL_TABLE_ENTRY_FORMAT_SHORT = '8sIhHBB'
SYMBOL_TABLE_ENTRY_FORMAT_LONG = 'IIIhHBB'
SYMBOL_TABLE_ENTRY_SIZE = calcsize(SYMBOL_TABLE_ENTRY_FORMAT_SHORT)
SYMBOL_TABLE_ENTRY_ZEROS = 0x0
SYMBOL_TABLE_ENTRY_SECTION = 1
SYMBOL_TABLE_ENTRY_ABSOLUTE = -1
SYMBOL_TABLE_ENTRY_TYPE = 0
SYMBOL_TABLE_ENTRY_CLASS = 2  # External (public) symbol.
SYMBOL_TABLE_ENTRY_CLASS_STATIC = 3
SYMBOL_TABLE_ENTRY_NUM_AUX = 0  # Number of auxiliary entries.

STRING_TABLE_OFFSET = 0x4  # Starting offset for the string table.
SIZE_FORMAT = 'I'
SIZE_LENGTH = calcsize(SIZE_FORMAT)

SIZE_SYMBOL_FORMAT_X64 = 'Q'
SIZE_SYMBOL_LENGTH_X64 = calcsize(SIZE_SYMBOL_FORMAT_X64)


def main():
    parser = argparse.ArgumentParser(
        description='Generate a COFF file for binary data.')
    parser.add_argument('--input', dest='input', help='Path of the input file.')
    parser.add_argument(
        '--output', dest='output', help='Name of the output file.')
    parser.add_argument(
        '--symbol_name',
        dest='symbol_name',
        help='Name of the symbol for the binary data')
    parser.add_argument(
        '--size_symbol_name',
        dest='size_name',
        help='Name of the symbol for the size of the binary data')
    parser.add_argument('--arch', dest='arch')
    parser.add_argument(
        '--executable', dest='executable', action='store_true', default=False)

    args = parser.parse_args()
    use_64_bit = args.arch in ['x64', 'arm64', 'riscv64']

    with open(args.input, 'rb') as f:
        section_data = f.read()

    # We need to calculate the following to determine the size of our buffer:
    #   1) Size of the data
    #   2) Total length of the symbol strings which are over 8 characters

    section_size = len(section_data)
    includes_size_name = (args.size_name != None)

    # Symbols on x86 are prefixed with '_'
    symbol_prefix = b'_' if args.arch == 'x86' else b''
    num_symbols = 1
    symbol_name = symbol_prefix + args.symbol_name.encode()
    size_symbol_name = None
    if (includes_size_name):
        size_symbol = args.size_name if args.size_name else args.symbol_name + "Size"
        size_symbol_name = symbol_prefix + size_symbol.encode()
        num_symbols += 1
    if (args.arch == 'x86'):
        feat_symbol_name = b'@feat.00'
        num_symbols += 1

    size_symbol_format = SIZE_SYMBOL_FORMAT_X64 if use_64_bit else SIZE_FORMAT
    size_symbol_size = SIZE_SYMBOL_LENGTH_X64 if use_64_bit else SIZE_LENGTH

    # The symbol table is directly after the data section
    symbol_table_ptr = (FILE_HEADER_SIZE + SECTION_HEADER_SIZE + section_size +
                        size_symbol_size)
    string_table_len = 0

    # Symbols longer than 8 characters have their string representations stored
    # in the string table.
    long_symbol_name = False
    long_size_symbol_name = False
    long_feat_symbol_name = False
    if (len(symbol_name) > SYMBOL_TABLE_ENTRY_SHORT_LEN):
        string_table_len += len(symbol_name) + 1
        long_symbol_name = True

    if (includes_size_name and
        (len(size_symbol_name) > SYMBOL_TABLE_ENTRY_SHORT_LEN)):
        string_table_len += len(size_symbol_name) + 1
        long_size_symbol_name = True

    if (args.arch == 'x86' and
        (len(feat_symbol_name) > SYMBOL_TABLE_ENTRY_SHORT_LEN)):
        string_table_len += len(feat_symbol_name) + 1
        long_feat_symbol_name = True

    # Create the buffer and start building.
    offset = 0
    buff = create_string_buffer(
        FILE_HEADER_SIZE + SECTION_HEADER_SIZE + section_size +
        num_symbols * SYMBOL_TABLE_ENTRY_SIZE + SIZE_LENGTH + size_symbol_size +
        string_table_len)

    FILE_HEADER_MAGIC = FILE_HEADER_MAGICS[args.arch]

    # Populate the file header. Basically constant except for the pointer to the
    # beginning of the symbol table.
    pack_into(FILE_HEADER_FORMAT, buff, offset, FILE_HEADER_MAGIC,
              FILE_HEADER_NUM_SECTIONS, FILE_HEADER_TIMESTAMP, symbol_table_ptr,
              num_symbols, FILE_HEADER_SIZE_OF_OPTIONAL, FILE_HEADER_FLAGS)
    offset += FILE_HEADER_SIZE

    section_name = SECTION_NAME_RDATA
    section_flags = SECTION_HEADER_DATA | SECTION_HEADER_ALIGN32 | SECTION_HEADER_READ
    if args.executable:
        section_name = SECTION_NAME_TEXT
        section_flags = SECTION_HEADER_TEXT | SECTION_HEADER_ALIGN32 | SECTION_HEADER_EXECUTE | SECTION_HEADER_READ

    # Populate the section header for a single section.
    pack_into(SECTION_HEADER_FORMAT, buff, offset, section_name, SECTION_PADDR,
              SECTION_VADDR, section_size + size_symbol_size,
              SECTION_RAW_DATA_PTR, SECTION_RELOCATION_PTR,
              SECTION_LINE_NUMS_PTR, SECTION_NUM_RELOCATION,
              SECTION_NUM_LINE_NUMS, section_flags)
    offset += SECTION_HEADER_SIZE

    # Copy the binary data.
    buff[offset:offset + section_size] = section_data
    offset += section_size

    # Append the size of the section.
    pack_into(size_symbol_format, buff, offset, section_size)
    offset += size_symbol_size

    # Build the symbol table. If a symbol name is 8 characters or less, it's
    # placed directly in the symbol table. If not, it's entered in the string
    # table immediately after the symbol table.

    string_table_offset = STRING_TABLE_OFFSET
    if long_symbol_name:
        pack_into(SYMBOL_TABLE_ENTRY_FORMAT_LONG, buff, offset,
                  SYMBOL_TABLE_ENTRY_ZEROS, string_table_offset, 0x0,
                  SYMBOL_TABLE_ENTRY_SECTION, SYMBOL_TABLE_ENTRY_TYPE,
                  SYMBOL_TABLE_ENTRY_CLASS, SYMBOL_TABLE_ENTRY_NUM_AUX)
        string_table_offset += len(symbol_name) + 1
    else:
        pack_into(SYMBOL_TABLE_ENTRY_FORMAT_SHORT, buff, offset, symbol_name,
                  0x0, SYMBOL_TABLE_ENTRY_SECTION, SYMBOL_TABLE_ENTRY_TYPE,
                  SYMBOL_TABLE_ENTRY_CLASS, SYMBOL_TABLE_ENTRY_NUM_AUX)
    offset += SYMBOL_TABLE_ENTRY_SIZE

    if includes_size_name:
        # The size symbol table entry actually contains the value for the size.
        if long_size_symbol_name:
            pack_into(SYMBOL_TABLE_ENTRY_FORMAT_LONG, buff, offset,
                      SYMBOL_TABLE_ENTRY_ZEROS, string_table_offset,
                      section_size, SYMBOL_TABLE_ENTRY_SECTION,
                      SYMBOL_TABLE_ENTRY_TYPE, SYMBOL_TABLE_ENTRY_CLASS,
                      SYMBOL_TABLE_ENTRY_NUM_AUX)
        else:
            pack_into(SYMBOL_TABLE_ENTRY_FORMAT_SHORT, buff, offset,
                      symbol_name, section_size, SYMBOL_TABLE_ENTRY_SECTION,
                      SYMBOL_TABLE_ENTRY_TYPE, SYMBOL_TABLE_ENTRY_CLASS,
                      SYMBOL_TABLE_ENTRY_NUM_AUX)
        offset += SYMBOL_TABLE_ENTRY_SIZE

    if args.arch == 'x86':
        FEATURE_SAFESEH = 0x1
        if long_feat_symbol_name:
            pack_into(SYMBOL_TABLE_ENTRY_FORMAT_LONG, buff, offset,
                      SYMBOL_TABLE_ENTRY_ZEROS, string_table_offset,
                      FEATURE_SAFESEH, SYMBOL_TABLE_ENTRY_ABSOLUTE,
                      SYMBOL_TABLE_ENTRY_TYPE, SYMBOL_TABLE_ENTRY_CLASS_STATIC,
                      SYMBOL_TABLE_ENTRY_NUM_AUX)
            string_table_offset += len(feat_symbol_name) + 1
        else:
            pack_into(SYMBOL_TABLE_ENTRY_FORMAT_SHORT, buff, offset,
                      feat_symbol_name, FEATURE_SAFESEH,
                      SYMBOL_TABLE_ENTRY_ABSOLUTE, SYMBOL_TABLE_ENTRY_TYPE,
                      SYMBOL_TABLE_ENTRY_CLASS_STATIC,
                      SYMBOL_TABLE_ENTRY_NUM_AUX)
            offset += SYMBOL_TABLE_ENTRY_SIZE

    pack_into(SIZE_FORMAT, buff, offset, string_table_len + SIZE_LENGTH)
    offset += SIZE_LENGTH

    # Populate the string table for any symbols longer than 8 characters.
    if long_symbol_name:
        symbol_len = len(symbol_name)
        buff[offset:offset + symbol_len] = symbol_name
        offset += symbol_len
        buff[offset] = b'\0'
        offset += 1

    if includes_size_name and long_size_symbol_name:
        symbol_len = len(size_symbol_name)
        buff[offset:offset + symbol_len] = size_symbol_name
        offset += symbol_len
        buff[offset] = b'\0'
        offset += 1

    if args.arch == 'x86' and long_feat_symbol_name:
        symbol_len = len(feat_symbol_name)
        buff[offset:offset + symbol_len] = feat_symbol_name
        offset += symbol_len
        buff[offset] = b'\0'
        offset += 1

    with open(args.output, 'wb') as f:
        f.write(buff.raw)


if __name__ == '__main__':
    main()