#!/usr/bin/env python3

import sys
import os
import re

FLASH_START_ADDR = 0x2000
FLASH_END_ADDR = 0x3fffc
NOP = bytes([0xc0, 0x46, 0xc0, 0x46])

WARNING = '\033[93m'
NORMAL  = '\033[0m'

unused_ranges = [[FLASH_START_ADDR, FLASH_END_ADDR]]
recovery_addr = 0
recovery_size = 0
recovery_bytes = bytes()
verbose = False

##############################
#    function definitions    #
##############################

def reversed_enumerate(L):
    for index in reversed(range(len(L))):
        yield index, L[index]

def break_range(rlo, rhi):
    global unused_ranges
    global verbose
    for i, a in reversed_enumerate(unused_ranges):
        if (rlo >= a[1]):
            break
        if (rhi >= a[0]):
            if (verbose):
                print('splitting range {:x}--{:x} using {:x}--{:x}'.format(a[0], a[1], rlo, rhi), end=' -> ')
            if (rlo - 4 > a[0]):
                unused_ranges[i] = [a[0], rlo - 4]
                if (rhi <= a[1]):
                    unused_ranges.insert(i+1, [rhi, a[1]])
                    if (verbose):
                        print('{:x}--{:x} and {:x}--{:x}'.format(a[0], rlo - 4, rhi, a[1]))
                else:
                    if (verbose):
                        print('{:x}--{:x}'.format(a[0], rlo - 4))
            else:
                if (rhi < a[1]):
                    unused_ranges[i] = [rhi, a[1]]
                    if (verbose):
                        print('{:x}--{:x}'.format(rhi, a[1]))
                else:
                    del unused_ranges[i]
                    if (verbose):
                        print('--')
            break

def print_range():
    global unused_ranges
    for a in unused_ranges:
        print('[{:x}, {:x}]'.format(a[0], a[1]))

def inunused(i):
    global unused_ranges
    for a in unused_ranges:
        if ((i >= a[0]) and (i <= a[1])):
            return((a[1] - a[0] + 4, a[1] - i + 4))
    return (0, 0)

def read_recovery_func():
    global recovery_addr
    global recovery_size
    global recovery_bytes
    with open(sys.argv[1], 'rb') as f:
        f.seek(recovery_addr - FLASH_START_ADDR)
        recovery_bytes = f.read(recovery_size)

def patch_bin():
    global unused_ranges
    global recovery_addr
    global recovery_size
    global recovery_bytes
    with open(sys.argv[1], 'r+b') as f:
        fsize = os.path.getsize(sys.argv[1])
        for i in range(0, fsize, 4):
            word = f.read(4)
            f.seek(i)
            (unused_block_size, space_left) = inunused(FLASH_START_ADDR + i)
            if (unused_block_size >= recovery_size + 4):
                if ((unused_block_size >= recovery_size) and (space_left <= recovery_size)
                    and (FLASH_START_ADDR + i + space_left != recovery_addr)):
                    print("Writing recovery block at address {:05x}...".format(i))
                    f.write(recovery_bytes[recovery_size - space_left:recovery_size - space_left + 4])
                else:
                    f.write(NOP)
                if (word != bytes([0, 0, 0, 0])):
                    print(WARNING + "WARNING: " + NORMAL + "overwriting actual data {:02x}{:02x}{:02x}{:02x} this probably shouldn't happen...".format(*[a for a in bytes(word)]))
            else:
                f.write(word)

##############################
#      main script code      #
##############################

if (len(sys.argv) < 2):
    exit(-1)

main_unit_name = re.match('([\w]+)\..*', sys.argv[1]).group(1)

if ((len(sys.argv) > 2) and (sys.argv[2] == '-v')):
    verbose = True

print('Using map file {:s} to patch {:s}...'.format(main_unit_name + '.map', main_unit_name + '.bin'))

prev_range_hi = 0
with open(main_unit_name + '.map', 'r') as f:
    lines = f.readlines()[1:]
    for l in lines:
        lparser = re.match('\s*([\w]+)\s+([\w]+)\s+([\w]+)\s+([\w]+)\s+(.*)', l)
        lma = int(lparser.group(2), 16)
        size = int(lparser.group(3), 16)
        symbol = lparser.group(5)
        if (re.match('_end = .', symbol)):
            break
        if (re.match('.text', symbol) or
            re.match('.*' + main_unit_name + '\.o.*', symbol) or
            re.match('\.filled', symbol) or
            re.match('\. \= ', symbol) or
            re.match('\$t', symbol) or
            re.match('\$d', symbol)):
            continue
        if (size == 0):
            continue
        if (re.match('emergency_recover', symbol)):
            recovery_addr = lma - (lma % 4)
            recovery_size = size
            if (verbose):
                print('Found recovery function at {:x}'.format(recovery_addr))
        if (verbose):
            print('found code block of size {:x} at {:x},'.format(size - ((lma + size) % 4) + 4, lma - (lma % 4)))
        break_range(lma - (lma % 4), (lma + size) - ((lma + size) % 4) + 4)
        prev_range_hi = lma + size

read_recovery_func()
patch_bin()

exit(0)
