#!/usr/bin/env python3
from absl import app
from absl import flags
from absl import logging
import json
import shutil
import struct
import os
import time
import datetime

FLAGS = flags.FLAGS

# Define command-line flags
flags.DEFINE_string('signer', None, 'Path to the unsigned user application firmware.')
flags.DEFINE_string('firmware', None, 'Path to the unsigned user application firmware.')
flags.DEFINE_string('bootloader', None, 'Path to the unsigned bootloader.')
flags.DEFINE_string('output_dir', None, 'Directory to save the signed files.')
flags.DEFINE_boolean('enable_swd', False, 'Enable SWD feature')
flags.DEFINE_boolean('boot_directly', False, 'Enable boot directly to user application')

# Ensure the flags are required
flags.mark_flag_as_required('output_dir')

def leftrotate(x, n):
    """Left rotate a 32-bit integer x by n bits."""
    return ((x << n) | (x >> (32 - n))) & 0xFFFFFFFF

def sha1_signer_python(data):
    """Python implementation matching sha1_signer binary exactly (based on reverse engineering)"""
    BLOCK_LEN = 64
    FIXED_SIZE = 0x1DF00  # sha1_signer processes exactly 0x1DF00 bytes
    
    # Create fixed-size buffer initialized with zeros (malloc + memset in C)
    buffer = bytearray(FIXED_SIZE)
    
    # Copy file data to buffer (fread in C)
    data_to_copy = min(len(data), FIXED_SIZE)
    buffer[:data_to_copy] = data[:data_to_copy]
    # Remaining bytes stay as 0
    
    # Custom initial hash values (confirmed from reverse engineering)
    hash_state = [0xef343ca7, 0x5c41a2de, 0x7ce9704f, 0x62312814, 0x4de11904]
    
    def sha1_compress(block, state):
        """Compress one 64-byte block (16 rounds only, confirmed from disassembly)"""
        a, b, c, d, e = state[0], state[1], state[2], state[3], state[4]
        
        # Process only 16 rounds (cmpl $0xf in disassembly)
        for i in range(16):
            # Load schedule using big-endian (confirmed from disassembly)
            schedule = (block[i * 4] << 24) | (block[i * 4 + 1] << 16) | \
                      (block[i * 4 + 2] << 8) | block[i * 4 + 3]
            
            # Round function with constant 0x7ebce2a7 (confirmed from disassembly)
            f = ((b & c) | (~b & d)) & 0xFFFFFFFF
            temp = (leftrotate(a, 5) + f + e + schedule + 0x7ebce2a7) & 0xFFFFFFFF
            e = d
            d = c
            c = leftrotate(b, 30)
            b = a
            a = temp
        
        # Update state
        state[0] = (state[0] + a) & 0xFFFFFFFF
        state[1] = (state[1] + b) & 0xFFFFFFFF
        state[2] = (state[2] + c) & 0xFFFFFFFF
        state[3] = (state[3] + d) & 0xFFFFFFFF
        state[4] = (state[4] + e) & 0xFFFFFFFF
    
    # Process complete 64-byte blocks from fixed-size buffer
    data_len = FIXED_SIZE
    offset = 0
    while data_len - offset >= BLOCK_LEN:
        sha1_compress(buffer[offset:offset + BLOCK_LEN], hash_state)
        offset += BLOCK_LEN
    
    # Handle remaining bytes and padding (standard SHA1 padding)
    block = bytearray(BLOCK_LEN)
    rem = data_len - offset
    if rem > 0:
        block[0:rem] = buffer[offset:offset + rem]
    
    # Add padding byte
    block[rem] = 0x80
    rem += 1
    
    # If not enough room for length, process this block and start a new one
    LENGTH_SIZE = 8
    if BLOCK_LEN - rem < LENGTH_SIZE:
        sha1_compress(bytes(block), hash_state)
        block = bytearray(BLOCK_LEN)
    
    # Add message length in bits (big-endian, standard SHA1 encoding)
    length_bits = data_len * 8
    for i in range(LENGTH_SIZE):
        block[BLOCK_LEN - 1 - i] = (length_bits >> (i * 8)) & 0xFF
    
    # Process final block
    sha1_compress(bytes(block), hash_state)
    
    return hash_state

def generate_metadata_python(file_path):
    """Generate metadata using Python implementation of sha1_signer."""
    with open(file_path, 'rb') as f:
        data = f.read()
    
    # Use exact clone of sha1_signer binary (based on reverse engineering)
    hash_values = sha1_signer_python(data)
    
    # Convert hash values to hex strings (same format as original)
    checksum = [f"{val:08x}" for val in hash_values]
    
    # Generate version number using UTC date format (YYYYMMDDHH)
    utc_now = datetime.datetime.utcnow()
    version_number = int(f"{utc_now.year:04d}{utc_now.month:02d}{utc_now.day:02d}{utc_now.hour:02d}")
    
    # Set copyright
    copyright_text = "Copyright (c) ifanr Inc, All Rights Reserved."
    
    return {
        "version_number": version_number,
        "checksum": checksum,
        "copyright": copyright_text
    }

def extract_release_id(file_path):
    """Extract the release ID from the binary."""
    with open(file_path, 'rb') as f:
        f.seek(0xA8)
        release_id = struct.unpack('<I', f.read(4))[0]
    return f"{release_id:x}"

def modify_firmware(firmware_bin, metadata, enable_swd = False, boot_directly = False):
    """Modifies the firmware binary based on the metadata."""
    if not enable_swd:
        # Insert 0xb10cdeb9 at 0x1DF10 to disable SWD
        firmware_bin = (
            firmware_bin[:0x1DF10]  # Everything before 0x1DF10
            + struct.pack("<I", 0xb10cdeb9)  # Insert 0xb10cdeb9 as 4-byte little-endian
            + firmware_bin[0x1DF14:]  # Everything after 0x1DF10
        )
    if boot_directly:
        # Insert 0xb007b007 at 0x1DF14 to boot directly to user application
        firmware_bin = (
            firmware_bin[:0x1DF14]  # Everything before 0x1DF14
            + struct.pack("<I", 0xb007b007)  # Insert 0xb007b007 as 4-byte little-endian
            + firmware_bin[0x1DF18:]  # Everything after 0x1DF14
        )

    version_number = metadata.get("version_number", 0)
    firmware_bin = (
        firmware_bin[:0x1DFF0]
        + struct.pack("<I", version_number)
        + firmware_bin[0x1DFF4:]
    )

    checksum = metadata.get("checksum", [0] * 5)
    for i, hexstr in enumerate(checksum):
        firmware_bin = (
            firmware_bin[: 0x1DFA0 + i * 4]
            + struct.pack("<I", int(hexstr, 16))
            + firmware_bin[0x1DFA4 + i * 4 :]
        )

    copyright = metadata.get("copyright", "").encode()
    copyright = copyright[:0x30]
    if len(copyright) < 0x30:
        copyright += b"\x00" * (0x30 - len(copyright))
    firmware_bin = firmware_bin[:0x1DFC0] + copyright + firmware_bin[0x1DFEF + 1 :]

    return firmware_bin

def modify_bootloader(firmware_bin, metadata):
    """Modifies the bootloader binary based on the metadata."""
    version_number = metadata.get('version_number', 0)
    firmware_bin = firmware_bin[:0x1ff4] + struct.pack('<I', version_number) + firmware_bin[0x1ff8:]

    flash_timestamp = metadata.get('flash_timestamp', 0)
    firmware_bin = firmware_bin[:0x1ff8] + struct.pack('<I', flash_timestamp) + firmware_bin[0x1ffc:]

    chip_id = metadata.get('chip_id', 0)
    firmware_bin = firmware_bin[:0x1ffe] + struct.pack('<B', chip_id) + firmware_bin[0x1fff:]

    copyright = metadata.get('copyright', '').encode()
    copyright = copyright[:0x30]
    if len(copyright) < 0x30:
        copyright += b'\x00' * (0x30 - len(copyright))
    firmware_bin = firmware_bin[:0x1fc0] + copyright + firmware_bin[0x1fef + 1:]

    return firmware_bin

def copy_associated_files(file_path, output_dir):
    """Copy the input file and associated map, txt, and elf files to the output directory."""
    base_name, _ = os.path.splitext(os.path.basename(file_path))
    for ext in ['.map', '.txt', '.elf']:
        src_file = file_path.replace('.bin', ext)
        if os.path.exists(src_file):
            shutil.copy(src_file, os.path.join(output_dir, os.path.basename(src_file)))

    shutil.copy(file_path, os.path.join(output_dir, os.path.basename(file_path)))

def main(argv):
    del argv  # Unused.

    if FLAGS.firmware:
        with open(FLAGS.firmware, 'rb') as f:
            firmware_bin = f.read()

        if len(firmware_bin) >= 0x1DF00:
            raise ValueError("Firmware is too big")

        firmware_bin += b"\xFF" * (0x1DF00 - len(firmware_bin))
        release_id = extract_release_id(FLAGS.firmware)
        unsigned_firmware_path = os.path.join(FLAGS.output_dir, f"firmware-unsigned-{release_id}.bin")
        with open(unsigned_firmware_path, 'wb') as f:
            f.write(firmware_bin)

        # Sign the firmware using Python implementation
        metadata = generate_metadata_python(unsigned_firmware_path)
        firmware_bin += b"\xFF" * 0x100
        signed_firmware_path = os.path.join(FLAGS.output_dir, f"firmware-{release_id}.bin")
        with open(signed_firmware_path, 'wb') as f:
            modified_firmware_bin = modify_firmware(firmware_bin, metadata, FLAGS.enable_swd, FLAGS.boot_directly)
            f.write(modified_firmware_bin)
            logging.info('Firmware %s written.', signed_firmware_path)

        # Write the metadata file
        metadata_path = os.path.join(FLAGS.output_dir, f"firmware-{release_id}.json")
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, indent=4)
            logging.info('Metadata %s written.', metadata_path)
        copy_associated_files(FLAGS.firmware, FLAGS.output_dir)

    if FLAGS.bootloader:
        # Sign the bootloader using Python implementation
        metadata = generate_metadata_python(FLAGS.bootloader)
        release_id = extract_release_id(FLAGS.bootloader)

        with open(FLAGS.bootloader, 'rb') as f:
            bootloader_bin = f.read()

        if len(bootloader_bin) > 0x1f00:
            raise ValueError('Bootloader is too big')
        bootloader_bin += b'\x00' * (0x1f00- len(bootloader_bin))
        bootloader_bin += b'\xff' * 0x10
        bootloader_bin += b'\x00' * 0xf0

        for chip_id in range(5):
            metadata['chip_id'] = chip_id
            signed_bootloader_path = os.path.join(FLAGS.output_dir, f"bootloader-chip{chip_id}-{release_id}.bin")
            with open(signed_bootloader_path, 'wb') as f:
                modified_bootloader_bin = modify_bootloader(bootloader_bin, metadata)
                f.write(modified_bootloader_bin)
                logging.info('Bootloader %s written.', signed_bootloader_path)

            # Write the metadata file
            metadata_path = os.path.join(FLAGS.output_dir, f"bootloader-{release_id}.json")
            with open(metadata_path, 'w') as f:
                json.dump(metadata, f, indent=4)
                logging.info('Metadata %s written.', metadata_path)
            copy_associated_files(FLAGS.bootloader, FLAGS.output_dir)

if __name__ == '__main__':
    app.run(main)