#!/usr/bin/env python3

#TODO - optimize it, should finally remove this packer

import os
import sys
import json
import shutil
import argparse
import copy
import errno
import re
import csv
import subprocess
import logging

armino_path = os.getenv("ARMINO_PATH")
project_dir = os.getenv("PROJECT_DIR")
project_name = os.getenv("PROJECT")
armino_soc = os.getenv("ARMINO_SOC")
build_path = os.path.realpath('.')
ota_tool = '%s/tools/env_tools/rtt_ota/ota-rbl/ota_packager_python.py'%(armino_path)
chip = 'bk7258'
firmware = 'app.bin'
image_name = 'all-app.bin'
file_op_list = list()
gen_file_list = set()
cpu1_bin  = f'{build_path}/../{armino_soc}_cp1/app.bin'
cpu2_bin  = f'{build_path}/../{armino_soc}_cp2/app.bin'

pack_boot_tools = '%s/tools/env_tools/beken_packager'%(armino_path)
armino_tool_part_table = '%s/tools/build_tools/part_table_tools/gen_bk7256partitions.py'%(armino_path)
partition_smode = '--smode'
partition_args = '16MB'
bootloader_json_inseq = '1,1,2,0,0,0'
bootloader_json_dirct = '%s/partition_bk7256_ota_a_new.json'%(pack_boot_tools)

def parse_args():
    description = '''Beken HID Downloader.'''
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('-c', '--chip',
                        default='bk7258',
                        help="chip type")
    parser.add_argument('-i', '--index', type=int,
                        default=0,
                        help="chip index")
    parser.add_argument('-b', '--boot',
                        help='specify boot file')
    parser.add_argument('-f', '--firmware',
                        help='specify firmware file')
    parser.add_argument('-d', '--dsp',
                        help='specify dsp file')
    parser.add_argument('-t', '--bt',
                        help='specify bt file')
    parser.add_argument('-n', '--image_name',
                        help='generated image name')
    parser.add_argument('-u', '--cpu1',
                        help='specify cup1 file')
    args = parser.parse_args()

    return args

#add new process
def run_cmd_with_ret(cmd):
    p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding='utf-8')
    out, err = p.communicate()
    print(out)
    if (p.returncode):
       print(err)
       exit(1)

#Size format conversion  string->int
def parse_int(v):
    for letter, multiplier in [("k", 1024), ("m", 1024 * 1024)]:
        if v.lower().endswith(letter):
            return parse_int(v[:-1]) * multiplier
    return int(v, 0)

#Size format conversion  int->string
def size_format(size, include_size):
    if include_size:
        for (val, suffix) in [(0x400, "K"), (0x100000, "M")]:
            if size % val == 0:
                return "%d%s" % (size // val, suffix)
    return "0x%x" % size

#Path check, if no, create
def ensure_directory(dir):
    if not os.path.exists(dir):
        try:
            os.makedirs(dir)
        except OSError as exc:
            if exc.errno != errno.EEXIST:
                raise

#physical address is translated into a logical address
def shrink_val(val, need_shrink):
    return int(val*32/34) if need_shrink else val

#use for debug
def show_gen_file_list():
    for file in gen_file_list:
        print("##$$>>:",file)

#move files from gen_file_list to new sub directory
def move_gen_files(sub_dir):
    global gen_file_list
    for file in gen_file_list:
        if os.path.exists(file.strip()):
            file_name = os.path.basename(file)
            file_dir = os.path.dirname(file)

            new_file = '%s/%s/%s'%(file_dir, sub_dir, file_name)
            nem_file_dir = os.path.dirname(new_file)
            ensure_directory(nem_file_dir)
            shutil.move(file, new_file)
            update_file_operation(file, new_file)
    gen_file_list = set()

#update file_op_list
def update_file_operation(src, new_src):
    for file_op_entry in file_op_list:
        if os.path.abspath(file_op_entry[0]) == os.path.abspath(src):
            file_op_entry[0] = new_src

#register file to file_op_list
def register_file_operation(src, dest, op):
    global file_op_list
    need_register = True
    if len(file_op_list) > 0:
        for s,d,o in file_op_list:
            if os.path.abspath(s) == os.path.abspath(src) and os.path.abspath(d) == os.path.abspath(dest) and o == op:
                need_register = False
                break
    if need_register:
        file_op_entry = [src, dest, op]
        file_op_list.append(file_op_entry)

#deal file to file_op_list
def deal_file_operation():
    global file_op_list
    for src,dest,op in file_op_list:
        if os.path.isfile(src):
            if op == 'move':
                ensure_directory(os.path.dirname(dest))
                shutil.move(src, dest)
            elif op == 'copy':
                ensure_directory(os.path.dirname(dest))
                shutil.copy(src, dest)
            elif op == 'rm':
                shutil.os.remove(src)
            else:
                pass
    file_op_list = list()

#get config json path
def get_config_json_path():
    json_path = None
    if os.path.exists('%s/%s/config/%s/configuration.json'%(armino_path, project_dir, chip)):
        json_path = '%s/%s/config/%s/configuration.json'%(armino_path, project_dir, chip)
    else:
        json_path = '%s/middleware/boards/%s/configuration.json'%(armino_path, chip)
    if json_path is None:
        raise 'config_json get failed!!!'
    register_file_operation('%s'%(json_path), "%s/encrypt/pack/configuration.json"%(build_path), 'copy')
    return json_path

#get partition csv path
def get_partition_csv_path():
    part_scv_path = None
    part_scv_name = '%s_partitions.csv'%(chip)
    if os.path.exists('%s/%s/config/%s/%s_partitions.csv'%(armino_path, project_dir, chip, chip)):
        part_scv_path = '%s/%s/config/%s/%s'%(armino_path, project_dir, chip,part_scv_name)
    elif os.path.exists('%s/middleware/boards/%s/partitions.csv'%(armino_path, chip)):
        part_scv_path = '%s/middleware/boards/%s/partitions.csv'%(armino_path, chip)
    else:
        raise RuntimeError('partition csv not exist,get csv failed!!!')

    return part_scv_path

#get the bootloader json for bootloader.bin inserting
def get_bootloader_json_for_inserting():
    csv_path = get_partition_csv_path()
    os.system('python3 %s %s %s --smode-inseq %s --flash-size %s --to-json %s'%(armino_tool_part_table, csv_path, partition_smode, bootloader_json_inseq, partition_args, bootloader_json_dirct))
    os.rename(bootloader_json_dirct,'%s/partition_bootloader.json'%(pack_boot_tools))

#get bootloader path
def get_bootloader_path(armino_path, chip):
    bootloader_no_insert_part_path = None
    if get_ab_position_independent_csv_status() == 'True':
        bootloader_no_insert_part_path = '%s/components/bk_libs/%s/bootloader/ab_bootloader'%(armino_path, chip)
    else:
        bootloader_no_insert_part_path = '%s/components/bk_libs/%s/bootloader/normal_bootloader'%(armino_path, chip)
    if os.path.exists('%s/partition_bootloader.json'%(pack_boot_tools)):
        print('project %s/%s in automatic_project list'%(project_name, armino_soc))
    else:
        print('project %s/%s not in automatic_project list'%(project_name, armino_soc))
        get_bootloader_json_for_inserting()

    os.system('%s/cmake_Gen_image genfile -injsonfile %s/config.json -infile %s/bootloader.bin -outfile %s/bootloader.bin -genjson %s/partition_bootloader.json'%(pack_boot_tools,pack_boot_tools, bootloader_no_insert_part_path,build_path, pack_boot_tools))

    bootloader_dest_path = '%s/bootloader.bin'%(build_path)

    return bootloader_dest_path

#get firmware path
def get_firmware_path(firmware):
    firmware_path = None
    if os.path.exists('%s/%s/config/%s/%s'%(armino_path, project_dir, chip, firmware)):
        firmware_path = '%s/%s/config/%s/%s'%(armino_path, project_dir, chip, firmware)
    else:
        if firmware == "bootloader.bin":
            firmware_path = get_bootloader_path(armino_path, chip)
        else:
            firmware_path = '%s/middleware/boards/%s/%s'%(armino_path, chip, firmware)

    if firmware == 'app1.bin':
        if os.path.exists(cpu1_bin):
            firmware_path = cpu1_bin;
        else:
            raise 'firmware: %s get failed!!!'%(firmware)

    if firmware == 'app2.bin':
        if os.path.exists(cpu2_bin):
            firmware_path = cpu2_bin;
        else:
            raise 'firmware: %s get failed!!!'%(firmware)

    if firmware_path is None:
        raise 'firmware: %s get failed!!!'%(firmware)

    if os.path.isfile(firmware_path):
        try:
            shutil.copy(firmware_path, "%s/%s"%(build_path, firmware))
            #register_file_operation("%s/%s"%(build_path, firmware), "%s/firmware/%s"%(build_path, firmware), 'move')
        except shutil.SameFileError:
            print('firmware  exsit')

    firmware_path = '%s/%s'%(build_path, firmware)
    register_file_operation('%s'%(firmware_path), "%s/encrypt/pack/%s"%(build_path, firmware), 'copy')
    return firmware_path

#copy bootloader.bin to components/bk_libs
def copy_bootloader_to_component_bklibs_dir(sour_boot_dir,dest_boot_dir):
    if not os.path.exists(dest_boot_dir):
        os.makedirs('%s/'%(dest_boot_dir))
    else:
        print("exsit",dest_boot_dir)
    shutil.copy('%s/%s/bootloader.bin'%(armino_path,sour_boot_dir),'%s/bootloader.bin'%(dest_boot_dir))

#copy bootloader bin-elf-asm to relating directory
def copy_bootloader_to_relevant_dir():
    project_bootloader_ab = "properties/modules/bootloader/aboot/arm_bootloader_ab/output"
    project_bootloader = "properties/modules/bootloader/aboot/arm_bootloader/output"
    if project_name != "ate_mini_code":
        if get_ab_position_independent_csv_status() == 'True':
            if os.path.exists('%s/%s/bootloader.bin'%(armino_path,project_bootloader_ab)):
                #copy bootloader bin to 'components/bk_libs/%s'%(chip)
                dest_boot_dir = '%s/components/bk_libs/%s/bootloader/ab_bootloader'%(armino_path, chip)
                copy_bootloader_to_component_bklibs_dir(project_bootloader_ab,dest_boot_dir)

                #copy bootloader elf-map-asm to build directory
                try:
                    shutil.copytree('%s/%s/'%(armino_path,project_bootloader_ab), './bootloader_out')
                except FileExistsError:
                    #raise RuntimeError("==========>./bootloader_out exist")
                    shutil.rmtree('./bootloader_out')
                    shutil.copytree('%s/%s/'%(armino_path,project_bootloader_ab), './bootloader_out')
        else:
            if os.path.exists('%s/%s/bootloader.bin'%(armino_path,project_bootloader)):
                print("PROJECT_DIR:",os.getenv("PROJECT_DIR"))
                #copy bootloader bin to 'components/bk_libs/%s'%(chip)
                dest_boot_dir = '%s/components/bk_libs/%s/bootloader/normal_bootloader'%(armino_path, chip)
                copy_bootloader_to_component_bklibs_dir(project_bootloader,dest_boot_dir)

                #copy bootloader and monitor elf-map-asm to build directory
                try:
                    shutil.copytree('%s/%s/'%(armino_path,project_bootloader), './bootloader_out')
                except FileExistsError:
                    #raise RuntimeError("==========>./bootloader_out exist")
                    shutil.rmtree('./bootloader_out')
                    shutil.copytree('%s/%s/'%(armino_path,project_bootloader), './bootloader_out')

#create image that only supports encryption 
#reads the size byte from the offset of src_file and overwrites dest_file to generate a new file
def create_only_support_enc(src_file, dest_file, offset, size):
    out_file = os.path.join(build_path, 'all_app_pack_enc_crc.bin')
    if not os.path.exists(out_file.strip()):
        shutil.copyfile(dest_file, out_file)
    with open(src_file, 'rb') as file1:
        src_data = bytearray(file1.read())
        data = src_data[offset:offset+size]
    with open(out_file, 'r+b') as f:
        f.seek(offset)
        f.write(data)
    gen_file_list.add(out_file)

#Package according to json file
def pack_from_config_json(cfg_json, nm):
    pack_tool = '%s/tools/env_tools/beken_packager/cmake_Gen_image'%(armino_path)
    json_file = './cfg_temp.json'
    infiles = ''
    if 'pack' in nm:
        outfile = os.path.join(build_path, '%s'%nm)
    else:
        nm_suffix = 'pack'
        outfile = os.path.join(build_path, '%s_%s'%(nm, nm_suffix))
    start_addr = cfg_json['section'][0]['start_addr']
    for sec in cfg_json['section']:
        infiles += '%s '%(sec['firmware'])
        start_addr = start_addr if parse_int(start_addr) <= parse_int(sec['start_addr']) else sec['start_addr']
    cfg_json_temp = copy.deepcopy(cfg_json)
    for sec in cfg_json_temp['section']:
        sec['start_addr'] = size_format(parse_int(sec['start_addr']) - parse_int(start_addr), False)
    cfg_json_temp = json.dumps(cfg_json_temp, sort_keys=False, indent=4)
    with open(json_file, 'w') as f:
        f.write(cfg_json_temp)
    if os.path.exists(pack_tool.strip()) and os.path.isfile(pack_tool.strip()):
        os.system('%s genfile -injsonfile %s -infile %s -outfile %s.bin'%(pack_tool, json_file, infiles, outfile))
    else:
        raise 'pack_tool path error!'
    os.remove(json_file)
    gen_file_list.add('%s.bin'%(outfile))
    return outfile

#add crc
def crc_from_config_json(cfg_json, nm):
    crc_tool = '%s/tools/env_tools/beken_packager/cmake_encrypt_crc'%(armino_path)
    nm_suffix = 'crc'
    outfile = os.path.join(build_path, '%s_%s'%(nm, nm_suffix))
    if os.path.exists(crc_tool.strip()) and os.path.isfile(crc_tool.strip()):
        os.system("%s -crc %s.bin"%(crc_tool, nm))
    else:
        raise 'crc_tool path error!'
    gen_file_list.add('%s.bin'%(outfile))
    return outfile

#Read the csv file
def read_csv_data(file_path):
    data = {}
    with open(file_path, 'r') as file:
        reader = csv.reader(file)
        next(reader)  # Skip header line
        for row in reader:
            data[row[0]] = row[1]
    return data

#get security ctrl status (bk7258 only support flash encrypt)
def get_security_status():
    if os.path.exists('%s/middleware/boards/%s/security_ctrl.csv'%(armino_path, chip)):
        security_ctrl_file = "%s/middleware/boards/%s/security_ctrl.csv"%(armino_path, chip)
        data = read_csv_data(security_ctrl_file)
        if data.get('security').lower() == 'false':
            return False
        for key,value in data.items():
            if key == project_name and value.lower() == 'true' :
                return True
    return False

#get ab position independent csv key value {TRUE :need position independent,FALSE :not need position independent }
def get_ab_pos_independent_value(ab_input_file):
    data = read_csv_data(ab_input_file)
    if data.get('pos_independent').lower() == 'true':
        output_value = 'True'
    else :
        output_value = 'False'

    return output_value

#get ab position independent csv status {TRUE :need position independent,FALSE :not need position independent }
def get_ab_position_independent_csv_status():
    if os.path.exists('%s/%s/config/%s/ab_position_independent.csv'%(armino_path, project_dir, chip)):
        ab_pos_independent_file = "%s/%s/config/%s/ab_position_independent.csv"%(armino_path, project_dir, chip)
        pos_independent_status = get_ab_pos_independent_value(ab_pos_independent_file)
    elif os.path.exists('%s/middleware/boards/%s/ab_position_independent.csv'%(armino_path, chip)):
        ab_pos_independent_file = "%s/middleware/boards/%s/ab_position_independent.csv"%(armino_path, chip)
        pos_independent_status = get_ab_pos_independent_value(ab_pos_independent_file)
    else:
        pos_independent_status = 'False'
    #print(' >>>>>>>>>>>>pos_independent_status: %s '%pos_independent_status)
    return pos_independent_status

# add ota head note into the all_app_crc.bin
def add_ota_head_into_all_app_image(dest_file, src_file):
    global write_data
    try :
        with open(dest_file , 'r+b') as dest_f:
            dest_f.seek(0x11000)
            try:
                with open(src_file,'rb') as src_f:
                    write_data = src_f.read()
            except FileNotFound:
                print(f'src_file not found {src_file}')
            dest_f.write(write_data)
    #os.remove('app_ab_crc.bin')
    except FileNotFound:
        print(f'dest_file not found {dest_file}')

#update the firmware version
#get firmware version from version_anti_rollback.csv
def update_firmware_version(img_type, nm):
    fm_version = 0
    if os.path.exists('%s/middleware/boards/%s/version_anti_rollback.csv'%(armino_path, chip)):
        version_csv_file = "%s/middleware/boards/%s/version_anti_rollback.csv"%(armino_path, chip)
        fm_version = int(read_csv_data(version_csv_file).get(img_type,0))
    else:
        fm_version = 0
    if img_type == 'bootloader': #current bootloader does not support version anti-rollback
        fm_version = 0

    with open('%s.bin'%(nm), 'rb') as file:
        content = bytearray(file.read())
    content[0x120] = fm_version
    with open('%s.bin'%(nm),'wb') as file:
        file.write(content)
    """
    with open('%s.bin'%(nm), 'r+b') as file:
        file.seek(0x120)
        file.write(fm_version)
    """

#find file by name,path search priority: project/config/chip > project > middleware/board/chip > tools/env_tools/beken_packager
def find_file_by_name(file_name):
    if os.path.exists('%s/%s/config/%s/%s'%(armino_path, project_dir, chip, file_name)):
        file = '%s/%s/%s/%s'%(armino_path, project_dir, chip, file_name)             #file from project/chip
    elif os.path.exists('%s/middleware/boards/%s/%s'%(armino_path, chip, file_name)):
        file = '%s/middleware/boards/%s/%s'%(armino_path, chip, file_name)           #file from middleware/board/chip
    elif os.path.exists('%s/middleware/boards/%s/csv/%s'%(armino_path, chip, file_name)):
        file = '%s/middleware/boards/%s/csv/%s'%(armino_path, chip, file_name)           #file from middleware/board/chip/csv
    else:
        file = None       #The default file used for test

    #print("<<<<<<debug file>>>>>>>",file)
    return file

#obtain the key pair for signature
#find the corresponding key based on the name configured in the csv
def get_private_key_for_sign(img_type):
    name = 'ecdsa384.der'
    if os.path.exists('%s/middleware/boards/%s/key_pair_selection.csv'%(armino_path, chip)):
        key_csv_file = "%s/middleware/boards/%s/key_pair_selection.csv"%(armino_path, chip)
        name = read_csv_data(key_csv_file).get(img_type)

    key = find_file_by_name(name)
    if key == None or not os.path.exists(key):
        key  = '%s/tools/env_tools/beken_packager/ecdsa384.der'%(armino_path)
    return key
#add sign
def sign_from_config_json(cfg_json, nm):
    sign_tool = '%s/tools/env_tools/beken_packager/SignTool_P384'%(armino_path)

    nm_suffix = 'sign'
    outfile = os.path.join(build_path, '%s_%s'%(nm, nm_suffix))
    size = 0
    
    for sec in cfg_json['section']:
        img_type = sec['partition']
        size += parse_int(sec['size'])
    if img_type == 'bootloader': # two levels of signatures
        prikey = get_private_key_for_sign('bootloader')
        update_firmware_version('bootloader', nm)
    else:
        prikey = get_private_key_for_sign('app')
        update_firmware_version('app', nm)

    if os.path.exists(sign_tool.strip()) and os.path.isfile(sign_tool.strip()):
        os.system("%s sign -infile %s.bin -prikey %s -outfile %s.bin -len %s"%(sign_tool, nm, prikey, outfile, hex(size)))
    else:
        raise 'sign_tool path error!'
    gen_file_list.add('%s.bin'%(outfile))
    return outfile

#obtain the aes key for flash encrypt
def get_aes_key_for_flash_encrypt():
    name = 'aes_encrypt_key.json'

    key = find_file_by_name(name)
    if key == None or not os.path.exists(key):
        key  = '%s/tools/env_tools/beken_packager/aes_encrypt_key.json'%(armino_path)# use default key
    return key

#get the aes key and iv for ota
def get_aes_key_iv_for_ota():
    key = get_aes_key_for_flash_encrypt()
    with open(key, "r") as f:
        key_json = json.load(f)
        for item in key_json:
            if item["name"] == "OTA_key_plaintext":
                aes_key = item["data"]
            if item["name"] == "OTA_IV_plaintext":
                aes_iv = item["data"]
    return aes_key, aes_iv

#obtain the aes key for flash encrypt
def get_aes_key_from_csv_for_flash_encrypt():
    name = 'security.csv'

    data = read_csv_data(find_file_by_name(name))
    aes_key = data.get('flash_aes_key')
    #print(">>>>>> debug flash encrypt key:",aes_key)
    if aes_key == None or aes_key == '':
        aes_key = '73c7bf397f2ad6bf4e7403a7b965dc5ce0645df039c2d69c814ffb403183fb18'
        print("the key for flash encryption not config in security.csv, use default key ",aes_key)

    return aes_key

#encrypt image
def enc_from_config_json(cfg_json, nm):
    enc_tool = '%s/tools/env_tools/beken_packager/beken_aes'%(armino_path)
    nm_suffix = 'enc'

    #key_file = get_aes_key_for_flash_encrypt()
    aes_key = get_aes_key_from_csv_for_flash_encrypt()
    start_addr = cfg_json['section'][0]['start_addr']
    for sec in cfg_json['section']:
        start_addr = start_addr if parse_int(start_addr) <= parse_int(sec['start_addr']) else sec['start_addr']

    outfile = os.path.join(build_path, '%s_%s'%(nm, nm_suffix))
    if os.path.exists(enc_tool.strip()) and os.path.isfile(enc_tool.strip()):
        os.system("%s encrypt -infile %s.bin -startaddress %s -keywords %s -outfile %s.bin"%(enc_tool, nm, start_addr, aes_key, outfile))
        #os.system("%s encrypt -infile %s.bin -startaddress %s -keyfile  %s -outfile %s.bin"%(enc_tool, nm, start_addr, key_file, outfile))
    else:
        raise 'enc_tool path error!'
    gen_file_list.add('%s.bin'%(outfile))
    gen_file_list.add('%s_crc.bin'%(outfile))
    return outfile

# add 0xff padding to binary file
def add_padding_to_binary(binary_file):
    with open(binary_file , 'ab')as f:
        f.write(bytes([0xff] * 34))  # 34 bytes align

def check_and_padding_binary(app_pack_secs, app_crc_name):
    last_firmware = app_pack_secs[-1]['firmware']
    last_firmware_size = os.path.getsize(last_firmware)
    last_partition_raw_size = parse_int(app_pack_secs[-1]['size'])
    last_firmware_aligned_size = (((last_firmware_size + 32 - 1) // 32) * 32) # 32 bytes align
    app_crc_name += '.bin'
    if last_firmware_aligned_size < last_partition_raw_size:
        add_padding_to_binary(app_crc_name) # if the last firmware not fill up the partition, add padding value.

#generate all app
def package_json_rebuild():
    global image_name
    header_path = "{}/tools/env_tools/rtt_ota/ota-rbl/".format(armino_path)
    app_pack_secs = list()

    json_file = get_config_json_path()
    with open(json_file, 'r') as local_json:
        config_json = json.load(local_json)
    for sec in config_json['section']:
        sec['firmware'] = get_firmware_path(sec['firmware'])
        sec['start_addr'] = size_format(shrink_val(parse_int(sec["start_addr"]), True), False)
        sec['size'] = size_format(shrink_val(parse_int(sec["size"]), True), True)
        if(sec['partition'] != 'bootloader'):
            app_pack_secs.append(sec)
    #gen all_app_pack.bin
    app_pack_name = pack_from_config_json(config_json, 'all_app')
    app_crc_name = crc_from_config_json(config_json, app_pack_name)
    check_and_padding_binary(app_pack_secs, app_crc_name)
    register_file_operation('%s.bin'%(app_pack_name), build_path, 'rm') #delete intermediate file
    
    src = '%s.bin'%(app_crc_name)
    dest = '%s/%s'%(build_path, image_name)
    op = 'move'
    register_file_operation(src, dest, op) #rename all-app.bin
    #gen app_pack.bin
    if len(app_pack_secs) > 0:
        config_json['count'] = len(app_pack_secs)
        config_json['section'] = app_pack_secs
    app_pack_name = pack_from_config_json(config_json, 'app')
    os.system('python3 %s -i %s.bin -o %s -g %s -ap %s -pjd %s packager'%(ota_tool, app_pack_name, 'app_pack.rbl', header_path, armino_path, project_dir))

    gen_file_list.add('%s/app_pack.rbl'%(build_path))

#generate all app
def package_json_enc_rebuild():
    global image_name
    header_path = "{}/tools/env_tools/rtt_ota/ota-rbl/".format(armino_path)
    app_pack_secs = list()

    json_file = get_config_json_path()
    with open(json_file, 'r') as local_json:
        config_json = json.load(local_json)
    for sec in config_json['section']:
        sec['firmware'] = get_firmware_path(sec['firmware'])
        sec['start_addr'] = size_format(shrink_val(parse_int(sec["start_addr"]), True), False)
        sec['size'] = size_format(shrink_val(parse_int(sec["size"]), True), True)
        if(sec['partition'] != 'bootloader'):
            app_pack_secs.append(sec)
    #gen all_app_pack.bin
    all_app_pack_name = pack_from_config_json(config_json, 'all_app')
    all_app_crc_name  = crc_from_config_json(config_json, all_app_pack_name)
    check_and_padding_binary(app_pack_secs, all_app_crc_name)
    all_app_enc_name  = enc_from_config_json(config_json, all_app_pack_name)
    #all_app_crc_name = crc_from_config_json(config_json, all_app_enc_name)
    create_only_support_enc('%s.bin'%(all_app_crc_name), '%s_crc.bin'%(all_app_enc_name), 0x110, 34)

    register_file_operation('%s.bin'%(all_app_crc_name), '%s/%s'%(build_path, image_name), 'move') #rename all-app.bin
    register_file_operation('%s.bin'%(all_app_pack_name), build_path, 'rm') #delete intermediate file
    #gen app_pack.bin
    if len(app_pack_secs) > 0:
        config_json['count'] = len(app_pack_secs)
        config_json['section'] = app_pack_secs
    app_pack_name = pack_from_config_json(config_json, 'app')
    app_enc_name = enc_from_config_json(config_json, app_pack_name)
    os.system('python3 %s -i %s.bin -o %s -g %s -ap %s -pjd %s packager'%(ota_tool, app_pack_name, 'app_pack.rbl', header_path, armino_path, project_dir))
    gen_file_list.add('%s/app_pack.rbl'%(build_path))

#generate all ab app
def package_json_ab_rebuild():
    image_nm = 'all_app_ab_crc.bin'
    ab_rbl_name = 'app_ab_crc.rbl'
    header_path = "{}/tools/env_tools/rtt_ota/ota-rbl/".format(armino_path)
    app_pack_secs = list()

    json_file = get_config_json_path()
    with open(json_file, 'r') as local_json:
        config_json = json.load(local_json)
    for sec in config_json['section']:
        sec['firmware'] = get_firmware_path(sec['firmware'])
        sec['start_addr'] = size_format(shrink_val(parse_int(sec["start_addr"]), True), False)
        sec['size'] = size_format(shrink_val(parse_int(sec["size"]), True), True)
        if(sec['partition'] != 'bootloader'):
            app_pack_secs.append(sec)
    #gen all_app_ab_crc.bin
    app_pack_name = pack_from_config_json(config_json, 'all_app')
    app_crc_name = crc_from_config_json(config_json, app_pack_name)
    check_and_padding_binary(app_pack_secs, app_crc_name)
    register_file_operation('%s.bin'%(app_pack_name), build_path, 'rm') #delete intermediate file

    src = '%s.bin'%(app_crc_name)
    dest = '%s/%s'%(build_path, image_nm)
    register_file_operation(src, dest, 'move') #rename all_app_ab_crc.bin
    if len(app_pack_secs) > 0:
        config_json['count'] = len(app_pack_secs)
        config_json['section'] = app_pack_secs
    app_pack_name = pack_from_config_json(config_json, 'app')
    register_file_operation('%s.bin'%(app_pack_name), build_path, 'rm') #delete intermediate file
    run_cmd_with_ret('python3 %s -i %s.bin -o %s -g %s -ap %s -soc %s -pjd %s packager'%(ota_tool, app_pack_name, 'app_ab.bin', header_path, armino_path, armino_soc, project_dir))
    app_ab = '%s/app_ab.bin'%(build_path)
    if os.path.isfile(app_ab):
        app_ab_crc = crc_from_config_json(config_json, 'app_ab')
        #register_file_operation('%s/encrypt/app_ab_crc.bin'%(build_path), '%s/app_ab_crc.bin'%(build_path), 'copy')
        add_ota_head_into_all_app_image(src, 'app_ab_crc.bin')
        #os.remove('app_ab_crc.bin')

#generate otp_efuse_config.json file for burn
def generate_otp_efuse_config_json():
    outfile =  os.path.join(build_path, 'otp_efuse_config.json')

    flash_aes_key = get_aes_key_from_csv_for_flash_encrypt()
    f = open(outfile, 'w+')
    otp_efuse_config = {
        "User_Operate_Enable":  "false",
        "Security_Ctrl_Enable": "true",
        "Security_Data_Enable": "true",

        "User_Operate":[],

        "Security_Ctrl":[{
            "secure_boot_enable":           "0,0,0",
            "secure_boot_debug_disable":    "0,1,0",
            "fast_boot_disable":            "0,2,0",
            "boot_mode":                    "0,3,0",
            "secure_boot_clock_select":     "0,4,0",
            "random_delay_enable":          "0,5,0",
            "direct_jump_enable":           "0,6,0",
            "boot_critical_error": "0,7,0",
            "attack_nmi_enable":            "2,4,0",
            "spi_to_ahb_disable":           "2,5,0",
            "auto_reset_enable[0]":         "2,6,0",
            "auto_reset_enable[1]":         "2,7,0",
            "flash_aes_enable":             "3,5,0",
            "spi_download_disable":         "3,6,0",
            "swd_disable":                  "3,7,0"
        }],

        "Security_Data":[]
    }
    data = {}

    otp_efuse_config["Security_Ctrl"][0].update({"flash_aes_enable":"3,5,1"})
    data["name"] = "flash_aes_key"
    data["mode"] = "write"
    data["permission"] = "WR" #TODO change to NA
    data["start_addr"] = "0x4B100460"
    data["last_valid_addr"] = "0x4B100480"
    data["byte_len"] = "0x20"
    data["data"] = flash_aes_key
    data["data_type"] = "hex"
    data["status"] = "true"
    otp_efuse_config["Security_Data"].append(data)

    json_str = json.dumps(otp_efuse_config, indent=4)
    with open('otp_efuse_config.json', 'w',newline="\n") as file:
        file.write(json_str)
    file.close()
    gen_file_list.add(outfile)

def mklittlefs_img():

    def make_img(img_size):
        mk_tool_path = "{}/components/littlefs/mkimg/mklittlefs".format(armino_path)
        preset_folder = "{}/littlefs".format(build_path)
        
        if not os.path.exists(preset_folder):
            os.makedirs(preset_folder)

        run_cmd_with_ret(f"{mk_tool_path} -c {preset_folder} -b 4096 -p 256 -s {img_size} {img_path}")

    def check_littlefs_enable():
        sdk_header_path = "{}/config/sdkconfig.h".format(build_path)
        with open(sdk_header_path, 'r', encoding='utf-8') as file:
            content = file.read()
        if "#define CONFIG_LITTLEFS 1" not in content:
            raise RuntimeError("CONFIG_LITTLEFS is not enable")

    def make_pack_bin_with_fs(img_offset):
        shutil.copy(all_app_bin_path, pack_bin_with_fs_path)
        with open(pack_bin_with_fs_path, 'r+b') as file:
            file.seek(0, os.SEEK_END)
            current_size = file.tell()
            # print(f"{current_size =}")
            
            if current_size < img_offset:
                bytes_to_fill = img_offset - current_size
                # print(f"{bytes_to_fill =}")
                fill_data = bytes([0xff] * bytes_to_fill)
                file.write(fill_data)

            with open(img_path, 'rb') as f2:
                fs_img_content = f2.read()

            file.write(fs_img_content)

    partition_csv_path = "{}/armino/partitions/_build/partitions.csv".format(build_path)
    img_path = "{}/littlefs.bin".format(build_path)
    all_app_bin_path = "{}/all-app.bin".format(build_path)
    pack_bin_with_fs_path = "{}/app_all_plus_filesys.bin".format(build_path)

    addr = 0
    size = 0
    found = 0

    with open(partition_csv_path, 'r', newline='', encoding='utf-8') as csvfile:
        dict_reader = csv.DictReader(csvfile)
        for row in dict_reader:
            if row['Name'] == 'littlefs':
                found = 1
                addr = int(row['Offset'], 16)
                size = parse_int(row['Size'])
                break

    if found == 0:
        print("not found littlefs partition in patitions.csv")
        return
    
    check_littlefs_enable()
    make_img(size)
    make_pack_bin_with_fs(addr)

def main():
    global chip
    global firmware
    global image_name
    args = parse_args()
    header_path = "{}/tools/env_tools/rtt_ota/ota-rbl/".format(armino_path)

    if args.chip is not None:
        chip = args.chip
    if args.firmware is not None:
        firmware = args.firmware
    if args.image_name is not None:
        image_name = args.image_name

    copy_bootloader_to_relevant_dir()

    if get_security_status() == True :
        generate_otp_efuse_config_json()
        package_json_enc_rebuild()
        #show_gen_file_list()
        move_gen_files('encrypt')
        deal_file_operation()
    else:
        if get_ab_position_independent_csv_status() == 'True':
            print('<<<<<<<<<<<<<<<<<<<<<<do_ab>>>>>>>>>>>>>>>>>>>')
            package_json_ab_rebuild()
            #show_gen_file_list()
            move_gen_files('encrypt')
            register_file_operation('%s/encrypt/app_ab_crc.bin'%(build_path), '%s/encrypt/app_ab_crc.rbl'%(build_path), 'move') #rename app_ab_crc.bin app_ab_crc.rbl
            deal_file_operation()
        else:
            package_json_rebuild()
            #show_gen_file_list()
            move_gen_files('encrypt')
            deal_file_operation()

    mklittlefs_img()


if __name__ == '__main__':
    try:
        main()
    except InputError as e:
        print(e, file=sys.stderr)
        sys.exit(2)