#!/bin/python

"""
This script patches the CRC checksums into an existing ELF file.

For this, it searches the follwoing sections:
	1) .text
	2) .data
	3) .ccmdata
	4) .vectors

All sections MUST be a multiple of 4 bytes long because the CRC calculation relies on whole 32 bit words.
The sections are excrated and the CRC is calculated for each section.

In the section .flashcrc, the script expects a single struct with the prototype:
struct flash_crcs {
	uint32_t start_magic;
	uint32_t crc_section_text;
	uint32_t crc_section_data;
	uint32_t crc_section_ccm_data;
	uint32_t crc_section_vectors;
	uint32_t end_magic;
};

It checks, if the start magic and end magic are set to the appropriate values and then patches in the CRC values of the sections.
The magic values checked for are: 0xA8BE53F9 and 0xFFA582FF
"""
from elftools.elf.elffile import ELFFile
from elftools.elf.segments import Segment
import elftools.elf.constants as elf_const
import sys
import crcmod
import crcmod.predefined
import struct

crc_calc = crcmod.predefined.mkCrcFun('crc-32-mpeg')

if len(sys.argv) < 2:
	print("Usage:", sys.argv[0], ' <elf file>')
	sys.exit(-1)

filename=sys.argv[1]

def section_calculate_crc(section):
	data = bytearray(section.data())
	be_data = bytearray([0 for k in range(0, len(data))])
	# Rearrange data, because the STM controller sees it as 32 bit little endian words
	for i in range(0, int(len(data)/4)):
		be_data[i*4+0] = data[i*4+3]
		be_data[i*4+1] = data[i*4+2]
		be_data[i*4+2] = data[i*4+1]
		be_data[i*4+3] = data[i*4+0]

	return crc_calc(be_data)

with open(filename, 'r+b') as f:
	elf = ELFFile(f)
	sections = {}
	sections['.text'] = elf.get_section_by_name('.text')
	sections['.data'] = elf.get_section_by_name('.data')
	sections['.ccmdata'] = elf.get_section_by_name('.ccmdata')
	sections['.vectors'] = elf.get_section_by_name('.vectors')

	for key, sec in sections.items():
		if sec is None:
			print("Error! Section", key, "not found in ELF file!")
			sys.exit(-1)
		print('Found section', key, 'Size:',
			  sec.data_size, 'Type:', sec['sh_type'])
		if sec['sh_type'] != 'SHT_PROGBITS':
			print('Error! Section must be of type SHT_PROGBITS')
			sys.exit(-1)
		if (sec.data_size % 4 != 0):
			print("Section", key, "has wrong size. Must be a multiple of 4 bytes!")
			sys.exit(-1)

	text_crc = section_calculate_crc(sections['.text'])
	print('CRC of .text section:', hex(text_crc))
	data_crc = section_calculate_crc(sections['.data'])
	print('CRC of .data section:', hex(data_crc))
	ccmdata_crc = section_calculate_crc(sections['.ccmdata'])
	print('CRC of .ccmdata section:', hex(ccmdata_crc))
	vextors_crc = section_calculate_crc(sections['.vectors'])
	print('CRC of .vectors section:', hex(vextors_crc))

	# Check the flashcrc section
	flashcrc_sec = elf.get_section_by_name('.flashcrc')
	if flashcrc_sec is None:
		print('Section for flash CRC missing!')
		sys.exit(-1)
	if flashcrc_sec.data_size != 6*4:
		print("Warning!!! .flashcrc section has wrong size:",flashcrc_sec.data_size)

	crc_sec_data = bytearray(flashcrc_sec.data())
	magic1 = struct.unpack('<I'*1, bytes(crc_sec_data[0:4]))[0]
	magic2 = struct.unpack('<I'*1, bytes(crc_sec_data[-4:]))[0]
	print("CRC section magic values:", hex(magic1), hex(magic2))
	if magic1 != 0xA8BE53F9 or magic2 != 0xFFA582FF:
		print("Wrong magics in CRC section. Data misalignment?")
		sys.exit(-2)

	crc_sec_offset = flashcrc_sec['sh_offset']
	print('CRC section ELF file offset:', hex(crc_sec_offset))
	
	crc_sec_data[4:8] = struct.pack('<I',text_crc)
	crc_sec_data[8:12] = struct.pack('<I',data_crc)
	crc_sec_data[12:16] = struct.pack('<I',ccmdata_crc)
	crc_sec_data[16:20] = struct.pack('<I',vextors_crc)
	f.seek(crc_sec_offset)
	f.write(crc_sec_data)
	print('CRCs patched successfully')