#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# This script was originally part of md380tools codebase by Travis Goodspeed

import argparse
import binascii
import struct
import sys

class TYTFW(object):
    def pad(self, align=512, byte=b'\xff'):
        pad_length = (align - len(self.app) % align) % align
        self.app += byte * pad_length

    def unwrap(self, img):
        header = struct.Struct(self.header_fmt)
        header = header.unpack(img[:256])

        self.start = header[6]
        app_len = header[7]
        self.app = self.crypt(img[256:256 + app_len])

        assert header[0].startswith(self.magic)
        assert header[1].startswith(self.jst)
        assert header[3].startswith(self.foo)
        assert header[4] == self.bar
        assert 0x8000000 <= header[6] < 0x8200000
        assert header[7] == len(img) - 512

    def crypt(self, data):
        return self.xor(data, self.key)

    @staticmethod
    def xor(a, b):
        # FIXME: optimized version
        out = bytearray()
        l = max(len(a), len(b))
        for i in range(l):
            out += bytes([a[i % len(a)] ^ b[i % len(b)]])
        return out

class MD380FW(TYTFW):
    # The stream cipher of MD-380 OEM firmware updates boils down
    # to a cyclic, static XOR key block, and here it is:
    key = (b'\x2e\xdf\x40\xb5\xbd\xda\x91\x35\x21\x42\xe3\xe2\x6d\xa9\x0b\x90'
           b'\x31\x30\x3a\xfa\x4f\x05\x74\x64\x0a\x29\x44\x7e\x60\x77\xad\x8c'
           b'\x9a\xe2\x63\xc4\x21\xfe\x3c\xf7\x93\xc2\xe1\x74\x16\x8c\xc9\x2a'
           b'\xed\x65\x68\x0c\x49\x86\xa3\xba\x61\x1c\x88\x5d\xc4\x49\x3c\xd2'
           b'\xee\x6b\x34\x0c\x1a\xa0\xa8\xb3\x58\x8a\x45\x11\xdf\x4f\x23\x2f'
           b'\xa4\xe4\xf6\x3b\x2c\x8c\x88\x2d\x9e\x9b\x67\xab\x1c\x80\xda\x29'
           b'\x53\x02\x1a\x54\x51\xca\xbf\xb1\x97\x22\x79\x81\x70\xfc\x00\xe9'
           b'\x81\x36\x4e\x4f\xa0\x1c\x0b\x07\xea\x2f\x49\x2f\x0f\x25\x71\xd7'
           b'\xf1\x30\x7d\x66\x6e\x83\x68\x38\x79\x13\xe3\x8c\x70\x9a\x4a\x9e'
           b'\xa9\xe2\xd6\x10\x4f\x40\x14\x8e\x6c\x5e\x96\xb2\x46\x3e\xe8\x25'
           b'\xef\x7c\xc5\x08\x18\xd4\x8b\x92\x26\xe3\xed\xfa\x88\x32\xe8\x97'
           b'\x47\x70\xf8\x46\xde\xff\x8b\x0c\x4d\xb3\xb6\xfc\x69\xd6\x27\x5b'
           b'\x76\x6f\x5b\x03\xf7\xc3\x11\x05\xc5\x1d\xfe\x92\x5f\xcb\xc2\x1c'
           b'\x81\x69\x1b\xb8\xf8\x62\x58\xc7\xb4\xb3\x11\xd5\x1f\xf2\x16\xc1'
           b'\xad\x8f\xa5\x1e\xb4\x5b\xe0\xda\x7f\x46\x7d\x1d\x9e\x6d\xc0\x74'
           b'\x7f\x54\xa6\x2f\x43\x6f\x64\x08\xca\xe8\x0f\x05\x10\x9c\x9d\x9f'
           b'\xbd\x67\x0c\x23\xf7\xa1\xe1\x59\x7b\xe8\xd4\x64\xec\x20\xca\xe9'
           b'\x6a\xb9\x03\x73\x67\x30\x95\x16\xb6\xd9\x19\x53\xe5\xdb\xa4\x3c'
           b'\xcd\x7c\xf9\xd8\x67\x9f\xfc\xc9\xe2\x8a\x6a\x2c\xf2\xed\xc8\xc1'
           b'\x6a\x20\x99\x4c\x0d\xad\xd4\x3b\xa1\x0e\x95\x88\x46\xb8\x13\xe1'
           b'\x06\x58\xd2\x07\xad\x5c\x1a\x74\xdb\xb5\xa7\x40\x57\xdb\xa2\x45'
           b'\xa6\x12\xd0\x82\xdd\xed\x0a\xbd\xb3\x10\xed\x6c\xda\x39\xd2\xd6'
           b'\x90\x82\x00\x76\x71\xe0\x21\xa0\x8f\xf0\xf3\x67\xc4\xf3\x40\xbd'
           b'\x47\x16\x10\xdc\x7e\xf8\x1d\xe5\x13\x66\x87\xc7\x4a\x69\xc9\x63'
           b'\x92\x82\xec\xee\x5a\x34\xfb\x96\x25\xc3\xb6\x68\xe1\x3c\x8a\x71'
           b'\x74\xb5\xc1\x23\x99\xd6\xf7\xfb\xea\x98\xcd\x61\x3d\x4d\xe1\xd0'
           b'\x34\xe1\xfd\x36\x10\x5f\x8e\x9e\xc6\xb6\x58\x0c\x55\xbe\x69\xa8'
           b'\x56\x76\x4b\x1f\xd5\x90\x7e\x47\x5f\x2f\x25\x02\x5c\xef\x00\x64'
           b'\xa0\x26\x9a\x18\x3c\x69\xc4\xff\x9a\x52\x41\x1b\xc9\x81\xc3\xac'
           b'\x15\xe1\x17\x98\xdb\x2c\x9c\x10\x9b\xb2\xf9\x71\x4f\x56\x0f\x68'
           b'\xfb\xd9\x2d\x5a\x86\x5b\x83\x03\xc8\x1e\xda\x5d\xe4\x8e\x82\xc3'
           b'\xd8\x7e\x8b\x56\x52\xb5\x38\xa0\xc6\xa9\xb0\x77\xbd\x8a\xf7\x24'
           b'\x70\x82\x1d\xc5\x95\x3c\xb5\xf0\x79\xa3\x89\x99\x4f\xec\x8c\x36'
           b'\xc7\xd6\x10\x20\xe3\x30\x39\x3d\x07\x9c\xb2\xdc\x4f\x94\x9e\xe0'
           b'\x24\xaa\xd2\x21\x12\x14\x41\x0f\xd4\x67\xb7\x99\xb1\xa3\xcb\x4d'
           b'\x0c\x70\x0f\xc0\x36\xa7\x89\x30\x86\x14\x67\x68\xac\x7b\xee\xe4'
           b'\x42\xd8\xb4\x36\xa4\xeb\x0f\xa8\x02\xf4\xcd\x23\xb3\xbc\x25\x4f'
           b'\xcc\xd4\xee\xfc\xf2\x21\x0f\xc1\x6c\x99\x37\xe2\x7c\x47\xce\x77'
           b'\xf0\x95\x2b\xcb\xf4\xca\x07\x03\x2a\xd2\x31\x00\xfd\x3e\x84\x86'
           b'\x32\x8b\x17\x9d\xbf\xa7\xb3\x37\xe1\xb1\x8a\x14\x69\x00\x25\xe3'
           b'\x56\x68\x9f\xaa\xa9\xb8\x11\x67\x75\x87\x4d\xf8\x36\x31\xcf\x38'
           b'\x63\x1c\xf0\x6b\x47\x40\x5d\xdc\x0c\xe6\xc8\xc4\x19\xaf\xdd\x6e'
           b'\x9e\xd9\x78\x99\x6c\xbe\x15\x1e\x0b\x9d\x88\xd2\x06\x9d\xee\xae'
           b'\x8a\x0f\xe3\x2d\x2f\xf4\xf5\xf6\x16\xbf\x59\xbb\x34\x5c\xdd\x61'
           b'\xed\x70\x1e\x61\xe5\xe3\xfb\x6e\x13\x9c\x49\x58\x17\x8b\xc8\x30'
           b'\xcd\xed\x56\xad\x22\xcb\x63\xce\x26\xc4\xa5\xc1\x63\x0d\x0d\x04'
           b'\x6e\xb6\xf9\xca\xbb\x2f\xab\xa0\xb5\x0a\xfa\x50\x0e\x02\x47\x05'
           b'\x54\x3d\xb3\xb1\xc6\xce\x8f\xac\x65\x7e\x15\x9e\x4e\xcc\x55\x9e'
           b'\x46\x32\x71\x9b\x97\xaa\x0d\xfb\x1b\x71\x02\x83\x96\x0b\x52\x77'
           b'\x48\x87\x61\x02\xc3\x04\x62\xd7\xfb\x74\x0f\x19\x9c\xa0\x9d\x79'
           b'\xa0\x6d\xef\x9e\x20\x5d\x0a\xc9\x6a\x58\xc9\xb9\x55\xad\xd1\xcc'
           b'\xd1\x54\xc8\x68\xc2\x76\xc2\x99\x0f\x2e\xfc\xfb\xf5\x92\xcd\xdb'
           b'\xa2\xed\xd9\x99\xff\x4f\x88\x50\xcd\x48\xb7\xb9\xf3\xf0\xad\x4d'
           b'\x16\x2a\x50\xaa\x6b\x2a\x98\x38\xc9\x35\x45\x0c\x03\xa8\xcd\x0d'
           b'\x74\x3c\x99\x55\xdb\x88\x70\xda\x6a\xc8\x34\x4d\x19\xdc\xcc\x42'
           b'\x40\x94\x61\x92\x65\x2a\xcd\xfd\x52\x10\x50\x14\x6b\xec\x85\x57'
           b'\x3f\xe2\x95\x9a\x5d\x11\xab\xad\x69\x60\xa8\x3b\x6f\x7a\x17\xf3'
           b'\x76\x17\x63\xe6\x59\x7e\x47\x30\xd2\x47\x87\xdb\xd8\x66\xde\x00'
           b'\x2b\x65\x37\x2f\x2d\xf1\x20\x11\xf3\x98\x7b\x4c\x9c\xd1\x76\xa7'
           b'\xe1\x3d\xbe\x6f\xee\x2c\xf0\x19\x70\x63\x51\x28\xf0\x1d\xbe\x52'
           b'\x5f\x4f\xe6\xde\xf2\x30\xb6\x50\x30\xf9\x15\x48\x49\xe9\xd2\xa8'
           b'\xa9\x8d\xda\xf5\xcd\x3e\xaf\x00\x55\xeb\x15\xc5\x5b\x19\x0f\x93'
           b'\x04\x27\x09\x6d\x54\xd7\x57\xb1\x47\x0a\xde\xf7\x1d\xcb\x11\x3c'
           b'\xf5\x8f\x20\x40\x9d\xbb\x6b\x2c\xa9\x67\x3d\x78\xc2\x62\xb7\x0c')

    def __init__(self, base_address=0x800c000):
        self.magic = b'OutSecurityBin'
        self.jst = b'JST51'
        self.foo = b'\x30\x02\x00\x30\x00\x40\x00\x47'
        self.bar = (b'\x01\x0d\x02\x03\x04\x05\x06\x07'
                    b'\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'
                    b'\x10\x11\x12\x13\x14\x15\x16\x17'
                    b'\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f'
                    b'\x20')
        self.start = base_address
        self.app = None
        self.footer = b'OutputBinDataEnd'
        self.header_fmt = '<16s7s9s16s33s47sLL120s'
        self.footer_fmt = '<240s16s'

    def wrap(self):
        bin = b''
        header = struct.Struct(self.header_fmt)
        footer = struct.Struct(self.footer_fmt)
        self.pad()
        app = self.crypt(self.app)
        bin += header.pack(
            self.magic, self.jst, b'\xff' * 9, self.foo,
            self.bar, b'\xff' * 47, self.start, len(app),
                                  b'\xff' * 120)
        bin += self.crypt(self.app)
        bin += footer.pack(b'\xff' * 240, self.footer)
        return bin


def main():
    def hex_int(x):
        return int(x, 0)

    parser = argparse.ArgumentParser(description='Wrap and unwrap MD-380 firmware')
    parser.add_argument('--wrap', '-w', dest='wrap', action='store_true',
                        default=False,
                        help='wrap app into firmware image')
    parser.add_argument('--unwrap', '-u', dest='unwrap', action='store_true',
                        default=False,
                        help='unwrap app from firmware image')
    parser.add_argument('--addr', '-a', dest='addr', type=hex_int,
                        default=0x800c000,
                        help='base address in flash')
    parser.add_argument('--offset', '-o', dest='offset', type=hex_int,
                        default=0,
                        help='offset to skip in app binary')
    parser.add_argument('input', nargs=1, help='input file')
    parser.add_argument('output', nargs=1, help='output file')
    args = parser.parse_args()

    if not (args.wrap ^ args.unwrap):
        sys.stderr.write('ERROR: --wrap or --unwrap?')
        sys.exit(5)

    print('DEBUG: reading "%s"' % args.input[0])
    with open(args.input[0], 'rb') as f:
        input = f.read()

    if args.wrap:
        if args.offset > 0:
            print('INFO: skipping 0x%x bytes in input file' % args.offset)
        
        md = MD380FW(args.addr)
        md.app = input[args.offset:]
        if len(md.app) == 0:
            sys.stderr.write('ERROR: seeking beyond end of input file\n')
            sys.exit(5)
        output = md.wrap()
        print('INFO: base address 0x{0:x}'.format(md.start))
        print('INFO: length 0x{0:x}'.format(len(md.app)))

    elif args.unwrap:
        md = MD380FW(args.addr)
        try:
            md.unwrap(input)
        except AssertionError:
            sys.stderr.write('WARNING: Funky header:\n')
            for i in range(0, 256, 16):
                hl = binascii.hexlify(input[i:i + 16])
                hl = ' '.join(hl[i:i + 2] for i in range(0, 32, 2))
                sys.stderr.write(hl + '\n')
            sys.stderr.write('Trying anyway.\n')
        output = md.app
        #print('INFO: base address 0x{0:x}'.format(md.start))
        print('INFO: length 0x{0:x}'.format(len(md.app)))

    print('DEBUG: writing "%s"' % args.output[0])
    with open(args.output[0], 'wb') as f:
        f.write(output)


if __name__ == "__main__":
    main()