#!/usr/bin/python3
#
# Copyright (c) 2015 Qualcomm Atheros, Inc.
# Copyright (c) 2018, The Linux Foundation. All rights reserved.
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

import struct
import ctypes
import os.path
import argparse
import json
import binascii
import hashlib
import tempfile
import subprocess
import logging
import sys
import shutil
import mailbox

MAX_BUF_LEN = 2000000

# the signature length also includes null byte and padding
ATH10K_BOARD_SIGNATURE = b"QCA-ATH10K-BOARD"
ATH10K_BOARD_SIGNATURE_LEN = 20

PADDING_MAGIC = 0x6d
DEFAULT_BD_API = 2
DEFAULT_BOARD_FILE = 'board-%d.bin' % DEFAULT_BD_API
DEFAULT_JSON_FILE = 'board-%d.json' % DEFAULT_BD_API
TYPE_LENGTH_SIZE = 8
BIN_SUFFIX = '.bin'

ATH10K_FIRMWARE_URL = 'https://github.com/kvalo/ath10k-firmware/commit'

ATH10K_BD_IE_BOARD = 0
ATH10K_BD_IE_BOARD_EXT = 1

ATH10K_BD_IE_BOARD_NAME = 0
ATH10K_BD_IE_BOARD_DATA = 1

ATH10K_BD_IE_BOARD_EXT_NAME = 0
ATH10K_BD_IE_BOARD_EXT_DATA = 1


def padding_needed(len):
    if len % 4 != 0:
        return 4 - len % 4
    return 0


def add_ie(buf, offset, id, value):
    length = len(value)
    padding_len = padding_needed(length)
    length = length + padding_len

    padding = ctypes.create_string_buffer(padding_len)

    for i in range(padding_len):
        struct.pack_into('<B', padding, i, PADDING_MAGIC)

    fmt = '<2i%ds%ds' % (len(value), padding_len)
    if not isinstance(value, bytes):
        value = value.encode()
    struct.pack_into(fmt, buf, offset, id, len(value), value, padding.raw)
    offset = offset + TYPE_LENGTH_SIZE + len(value) + padding_len

    return offset


def xclip(msg):
    p = subprocess.Popen(['xclip', '-selection', 'clipboard'],
                         stdin=subprocess.PIPE)
    p.communicate(msg.encode())


# to workaround annoying python feature of returning negative hex values
def hex32(val):
    return val & 0xffffffff


# match with kernel crc32_le(0, buf, buf_len) implementation
def _crc32(buf):
    return hex32(~(hex32(binascii.crc32(buf, 0xffffffff))))


def pretty_array_str(array):
    return '\',\''.join(array)


class BoardName():
    @staticmethod
    def parse_ie(buf, offset, length):
        self = BoardName()
        fmt = '<%ds' % length
        (name, ) = struct.unpack_from(fmt, buf, offset)
        self.name = name.decode()

        logging.debug('BoardName.parse_ie(): offset %d length %d self %s' %
                      (offset, length, self))

        return self

    def add_to_buf(self, ebdf, buf, offset):
        ie_id = ATH10K_BD_IE_BOARD_NAME
        if ebdf:
            ie_id = ATH10K_BD_IE_BOARD_EXT_NAME

        return add_ie(buf, offset, ie_id, self.name.encode('utf-8'))

    def __eq__(self, other):
        return self.name == other.name

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return 'BoardName(%s)' % self.name

    def __init__(self, name=None):
        self.name = name


class BoardData():
    @staticmethod
    def parse_ie(buf, offset, length):
        self = BoardData()
        fmt = '<%ds' % length
        (self.data, ) = struct.unpack_from(fmt, buf, offset)

        logging.debug('BoardData.parse_ie(): offset %d length %d self %s' %
                      (offset, length, self))

        return self

    def add_to_buf(self, ebdf, buf, offset):
        ie_id = ATH10K_BD_IE_BOARD_DATA
        if ebdf:
            ie_id = ATH10K_BD_IE_BOARD_EXT_DATA

        return add_ie(buf, offset, ie_id, self.data)

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        if self.data is not None:
            s = '%d B' % (len(self.data))
        else:
            s = 'n/a'

        return 'BoardData(%s)' % (s)

    def __init__(self, data=None):
        self.data = data


class Board():
    @staticmethod
    def parse_ie(buf, offset, length):
        logging.debug('Board.parse_ie(): offset %d length %d' % (offset, length))

        self = Board()

        # looping board IEs
        while length > 0:
            (ie_id, ie_len) = struct.unpack_from('<2i', buf, offset)

            logging.debug('Board.parse_ie(): found ie_id %d ie_len %d offset %d length %d' %
                          (ie_id, ie_len, offset, length))

            if TYPE_LENGTH_SIZE + ie_len > length:
                raise Exception('Error: ie_len too big (%d > %d)' % (ie_len,
                                                                     length))

            offset += TYPE_LENGTH_SIZE
            length -= TYPE_LENGTH_SIZE

            # FIXME: create separate ExtenderBoard() class so that we
            # don't need these "or" hacks here. Maybe add a common
            # abstract class to avoid code duplication?
            if ie_id == ATH10K_BD_IE_BOARD_NAME or ie_id == ATH10K_BD_IE_BOARD_EXT_NAME:
                self.names.append(BoardName.parse_ie(buf, offset, ie_len))
            elif ie_id == ATH10K_BD_IE_BOARD_DATA or ie_id == ATH10K_BD_IE_BOARD_EXT_DATA:
                self.data = BoardData.parse_ie(buf, offset, ie_len)

            offset += ie_len + padding_needed(ie_len)
            length -= ie_len + padding_needed(ie_len)

        return self

    def add_to_buf(self, buf, offset):
        # store position ie header of this element
        ie_offset = offset

        offset += TYPE_LENGTH_SIZE

        ie_id = ATH10K_BD_IE_BOARD

        if self.ebdf:
            ie_id = ATH10K_BD_IE_BOARD_EXT

        for name in self.names:
            offset = name.add_to_buf(self.ebdf, buf, offset)

        offset = self.data.add_to_buf(self.ebdf, buf, offset)

        # write ie header as now we know the full length
        ie_len = offset - ie_offset - TYPE_LENGTH_SIZE
        struct.pack_into('<2i', buf, ie_offset, ie_id, ie_len)

        return offset

    def get_names_as_str(self):
        names = []
        for boardname in self.names:
            names.append(boardname.name)

        return names

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        names = []
        for boardname in self.names:
            names.append(str(boardname))

        return 'Board(%s, %s)' % (','.join(names), self.data)

    def __init__(self):
        self.data = None
        self.names = []
        self.ebdf = False


class BoardContainer:

    def find_board(self, name):
        for board in self.boards:
            for boardname in board.names:
                if name == boardname.name:
                    return board

    def add_board(self, data, names):
        boardnames = []
        ebdf = False

        for name in names:
            b = self.find_board(name)
            if b:
                self.boards.remove(b)

            if "bmi-eboard-id" in name:
                ebdf = True
            boardnames.append(BoardName(name))

        board = Board()
        board.ebdf = ebdf
        board.data = BoardData(data)
        board.names = boardnames

        self.boards.append(board)

    @staticmethod
    def open_json(filename):
        self = BoardContainer()

        if not os.path.exists(filename):
            print('mapping file %s not found' % (filename))
            return

        f = open(filename, 'r')
        mapping = json.loads(f.read())
        f.close()

        for b in mapping:
            board_filename = b['data']
            f = open(board_filename, 'rb')
            data = f.read()
            f.close()

            self.add_board(data, b['names'])

        return self

    def validate(self):
        allnames = []

        for board in self.boards:
            for name in board.names:
                if name in allnames:
                    # TODO: Find a better way to report problems,
                    # maybe return a list of strings? Or use an
                    # exception?
                    print('Warning: duplicate board name: %s' % (name.name))
                    return

                allnames.append(name)

    def _add_signature(self, buf, offset):
        signature = ATH10K_BOARD_SIGNATURE + b'\0'
        length = len(signature)
        pad_len = padding_needed(length)
        length = length + pad_len

        padding = ctypes.create_string_buffer(pad_len)

        for i in range(pad_len):
            struct.pack_into('<B', padding, i, PADDING_MAGIC)

        fmt = '<%ds%ds' % (len(signature), pad_len)
        struct.pack_into(fmt, buf, offset, signature, padding.raw)
        offset += length

        # make sure ATH10K_BOARD_SIGNATURE_LEN is correct
        assert ATH10K_BOARD_SIGNATURE_LEN == length

        return offset

    @staticmethod
    def open(name):
        self = BoardContainer()

        f = open(name, 'rb')
        buf = f.read()
        f.close()
        buf_len = len(buf)

        logging.debug('BoardContainer.open(): name %s' % (name))

        offset = 0

        fmt = '<%dsb' % (len(ATH10K_BOARD_SIGNATURE))
        (signature, null) = struct.unpack_from(fmt, buf, offset)

        if signature != ATH10K_BOARD_SIGNATURE or null != 0:
            print("invalid signature found in %s" % name)
            return 1

        offset += ATH10K_BOARD_SIGNATURE_LEN

        # looping main IEs
        while offset < buf_len:
            (ie_id, ie_len) = struct.unpack_from('<2i', buf, offset)
            logging.debug('BoardContainer.open(): found offset %d ie_id %d ie_len %d' %
                          (offset, ie_id, ie_len))

            offset += TYPE_LENGTH_SIZE

            if offset + ie_len > buf_len:
                print('Error: Buffer too short (%d + %d > %d)' % (offset,
                                                                  ie_len,
                                                                  buf_len))
                return 1

            # FIXME: create separate ExtenderBoard() class and
            # self.extender_boards list
            if ie_id == ATH10K_BD_IE_BOARD or ie_id == ATH10K_BD_IE_BOARD_EXT:
                self.boards.append(Board.parse_ie(buf, offset, ie_len))

            offset += ie_len + padding_needed(ie_len)

        self.validate()

        return self

    def write(self, name):
        (buf, buf_len) = self.get_bin()

        fd = open(name, 'wb')
        fd.write(buf.raw[0:buf_len])
        fd.close()

        self.validate()

        print("board binary file '%s' is created" % name)

    def get_bin(self):
        buf = ctypes.create_string_buffer(MAX_BUF_LEN)
        offset = 0

        offset = self._add_signature(buf, offset)

        for board in self.boards:
            offset = board.add_to_buf(buf, offset)

        # returns buffer and it's length
        return buf, offset

    def get_summary(self, sort=False):
        (buf, buf_len) = self.get_bin()

        s = ''

        s += 'FileSize: %d\n' % (buf_len)
        s += 'FileCRC32: %08x\n' % (_crc32(buf[0:buf_len]))
        s += 'FileMD5: %s\n' % (hashlib.md5(buf[0:buf_len]).hexdigest())

        boards = self.boards

        if sort:
            boards = sorted(boards, key=lambda board: board.get_names_as_str())

        index = 0
        for board in boards:
            if not sort:
                index_s = '[%d]' % (index)
            else:
                index_s = ''

            s += 'BoardNames%s: \'%s\'\n' % (index_s, pretty_array_str(board.get_names_as_str()))
            s += 'BoardLength%s: %d\n' % (index_s, len(board.data.data))
            s += 'BoardCRC32%s: %08x\n' % (index_s, _crc32(board.data.data))
            s += 'BoardMD5%s: %s\n' % (index_s, hashlib.md5(board.data.data).hexdigest())
            index += 1

        return s

    def __init__(self):
        self.boards = []


def cmd_extract(args):
    cont = BoardContainer().open(args.extract)

    mapping = []

    for board in cont.boards:
        filename = board.names[0].name + '.bin'

        b = {}
        b['names'] = board.get_names_as_str()
        b['data'] = filename
        mapping.append(b)

        f = open(filename, 'wb')
        f.write(board.data.data)
        f.close()

        print("%s created size: %d" % (filename, len(board.data.data)))

    filename = DEFAULT_JSON_FILE
    f = open(filename, 'w')
    f.write(json.dumps(mapping, indent=4))
    f.close()

    print("%s created" % (filename))


def cmd_info(args):
    filename = args.info

    cont = BoardContainer().open(filename)

    print(cont.get_summary())


def cmd_diff(args):
    if args.diff:
        filename1 = args.diff[0]
        filename2 = args.diff[1]
    else:
        filename1 = args.diffstat[0]
        filename2 = args.diffstat[1]

    print(diff_boardfiles(filename1, filename2, args.diff))


def diff_boardfiles(filename1, filename2, diff):
    result = ''

    container1 = BoardContainer().open(filename1)
    (temp1_fd, temp1_pathname) = tempfile.mkstemp()
    os.write(temp1_fd, container1.get_summary(sort=True).encode())

    container2 = BoardContainer().open(filename2)
    (temp2_fd, temp2_pathname) = tempfile.mkstemp()
    os.write(temp2_fd, container2.get_summary(sort=True).encode())

    # this function is used both with --diff and --diffstat
    if diff:
        # '--less-mode' and '--auto-page' would be nice when running on
        # terminal but don't know how to get the control character
        # through. For terminal detection sys.stdout.isatty() can be used.
        cmd = ['wdiff', temp1_pathname, temp2_pathname]

        # wdiff is braindead and returns 1 in a succesfull case
        try:
            output = subprocess.check_output(cmd)
        except subprocess.CalledProcessError as e:
            if e.returncode == 1:
                output = e.output
            else:
                print('Failed to run wdiff: %d\n%s' % (e.returncode, e.output))
                return 1
        except OSError as e:
            print('Failed to run wdiff: %s' % (e))
            return 1

        result += '%s\n' % (output.decode())

    # create simple statistics about changes in board images

    new_boards = {}
    deleted_boards = {}
    changed_boards = {}

    for board in container2.boards:
        # convert the list to a string
        s = pretty_array_str(board.get_names_as_str())
        new_boards[s] = board

    for board in container1.boards:
        # convert the list to a string
        names = pretty_array_str(board.get_names_as_str())

        if names not in new_boards:
            # board image has been deleted
            deleted_boards[names] = board
            continue

        board2 = new_boards[names]
        del new_boards[names]

        if board.data.data == board2.data.data:
            # board image hasn't changed
            continue

        # board image has changed
        changed_boards[names] = board2

    result += 'New:\n%s\n\n' % ('\n'.join(new_boards.keys()))
    result += 'Changed:\n%s\n\n' % ('\n'.join(changed_boards.keys()))
    result += 'Deleted:\n%s\n' % ('\n'.join(deleted_boards.keys()))

    result += '%d board image(s) added, %d changed, %d deleted, %d in total' % (len(new_boards),
                                                                                len(changed_boards),
                                                                                len(deleted_boards),
                                                                                len(container2.boards))

    os.close(temp1_fd)
    os.close(temp2_fd)

    return result


def cmd_create(args):
    mapping_file = args.create

    if args.output:
        output = args.output
    else:
        output = DEFAULT_BOARD_FILE

    cont = BoardContainer.open_json(mapping_file)
    cont.write(output)


def cmd_add_board(args):
    if len(args.add_board) < 3:
        print('error: --add-board requires 3 or more arguments, only %d given' % (len(args.add_board)))
        sys.exit(1)

    board_filename = args.add_board[0]
    new_filename = args.add_board[1]
    new_names = args.add_board[2:]

    f = open(new_filename, 'rb')
    new_data = f.read()
    f.close()

    # copy the original file for diff
    (temp_fd, temp_pathname) = tempfile.mkstemp()
    shutil.copyfile(board_filename, temp_pathname)

    container = BoardContainer.open(board_filename)
    container.add_board(new_data, new_names)
    container.write(board_filename)

    print(diff_boardfiles(temp_pathname, board_filename, False))

    os.remove(temp_pathname)


def git_get_head_id():
    return subprocess.check_output(['git', 'log', '--format=format:%H', '-1'], text=True)


def cmd_add_mbox(args):
    board_filename = args.add_mbox[0]
    mbox_filename = args.add_mbox[1]

    mbox = mailbox.mbox(mbox_filename)

    for msg in mbox:
        board_files = {}

        for part in msg.walk():
            filename = part.get_filename()

            if not filename:
                # not an attachment
                continue

            if not filename.endswith(BIN_SUFFIX):
                continue

            name = filename[:-len(BIN_SUFFIX)]
            board_files[name] = part.get_payload(decode=True)

        print('Found mail "%s" with %d board files' % (msg['Subject'],
                                                       len(board_files)))

        # copy the original file for diff
        (temp_fd, temp_pathname) = tempfile.mkstemp()
        shutil.copyfile(board_filename, temp_pathname)

        container = BoardContainer.open(board_filename)
        for name, data in board_files.items():
            names = [name]
            container.add_board(data, names)

        container.write(board_filename)

        applied_msg = ''

        if args.commit:
            cmd = ['git', 'commit', '--quiet', '-m', msg['Subject'], board_filename]
            subprocess.check_call(cmd)

            applied_msg += 'Thanks, added to %s:\n\n' % (board_filename)

        applied_msg += diff_boardfiles(temp_pathname, board_filename, False)
        applied_msg += '\n\n'

        if args.commit:
            applied_msg += '%s/%s\n\n' % (ATH10K_FIRMWARE_URL, git_get_head_id())

        os.remove(temp_pathname)

        print('----------------------------------------------')
        print(applied_msg)
        print('----------------------------------------------')

        xclip(applied_msg)


def main():
    description = '''ath10k board-N.bin files manegement tool

ath10k-bdencoder is for creating (--create), listing (--info) and
comparing (--diff, --diffstat) ath10k board-N.bin files. The
board-N.bin is a container format which can have unlimited number of
actual board images ("board files"), each image containing one or
names which ath10k uses to find the correct image.

For creating board files you need a mapping file in JSON which
contains the names and filenames for the actual binary:

[
        {"names": ["AAA1", "AAAA2"], "data": "A.bin"},
        {"names": ["B"], "data": "B.bin"},
        {"names": ["C"], "data": "C.bin"},
]

In this example the board-N.bin will contain three board files which
are read from files named A.bin (using names AAA1 and AAAA2 in the
board-N.bin file), B.bin (using name B) and C.bin (using name C). You
can use --extract switch to see examples from real board-N.bin files.
'''

    parser = argparse.ArgumentParser(description=description,
                                     formatter_class=argparse.RawTextHelpFormatter)

    cmd_group = parser.add_mutually_exclusive_group(required=True)
    cmd_group.add_argument("-c", "--create", metavar='JSON_FILE',
                           help='create board-N.bin from a mapping file in JSON format')
    cmd_group.add_argument("-e", "--extract", metavar='BOARD_FILE',
                           help='extract board-N.bin file to a JSON mapping file and individual board files, compatible with the format used with --create command')
    cmd_group.add_argument("-i", "--info", metavar='BOARD_FILE',
                           help='show all details about a board-N.bin file')
    cmd_group.add_argument('-d', '--diff', metavar='BOARD_FILE', nargs=2,
                           help='show differences between two board-N.bin files')
    cmd_group.add_argument('-D', '--diffstat', metavar='BOARD_FILE', nargs=2,
                           help='show a summary of differences between two board-N.bin files')
    cmd_group.add_argument('-a', '--add-board', metavar='NAME', nargs='+',
                           help='add a board file to an existing board-N.bin, first argument is the filename of board-N.bin to add to, second is the filename board file (board.bin) to add and then followed by one or more arguments are names used in board-N.bin')

    cmd_group.add_argument('-A', '--add-mbox', metavar='FILENAME', nargs=2,
                           help='FIXME')

    parser.add_argument('-C', '--commit', action='store_true',
                        help='commit changes to a git repository')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='enable verbose (debug) messages')
    parser.add_argument("-o", "--output", metavar="BOARD_FILE",
                        help='name of the output file, otherwise the default is: %s' %
                        (DEFAULT_BOARD_FILE))

    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)

    if args.create:
        return cmd_create(args)
    elif args.extract:
        return cmd_extract(args)
    elif args.info:
        return cmd_info(args)
    elif args.diff:
        return cmd_diff(args)
    elif args.diffstat:
        return cmd_diff(args)
    elif args.add_board:
        return cmd_add_board(args)
    elif args.add_mbox:
        return cmd_add_mbox(args)

if __name__ == "__main__":
    main()