1*deb4d487SHisping Lin#!/usr/bin/env python3 2*deb4d487SHisping Lin# 3*deb4d487SHisping Lin# Copyright 2023, Rockchip Electronics Co., Ltd 4*deb4d487SHisping Lin# hisping lin, <hisping.lin@rock-chips.com> 5*deb4d487SHisping Lin# 6*deb4d487SHisping Lin# SPDX-License-Identifier: GPL-2.0+ 7*deb4d487SHisping Lin# 8*deb4d487SHisping Lin 9*deb4d487SHisping Lindef get_args(): 10*deb4d487SHisping Lin import argparse 11*deb4d487SHisping Lin 12*deb4d487SHisping Lin parser = argparse.ArgumentParser() 13*deb4d487SHisping Lin 14*deb4d487SHisping Lin parser.add_argument('--out', required=True, \ 15*deb4d487SHisping Lin help='Name of image for the TA') 16*deb4d487SHisping Lin 17*deb4d487SHisping Lin parser.add_argument('--tadir', required=True, help='Path of TA director') 18*deb4d487SHisping Lin 19*deb4d487SHisping Lin return parser.parse_args() 20*deb4d487SHisping Lin 21*deb4d487SHisping Lindef main(): 22*deb4d487SHisping Lin import os 23*deb4d487SHisping Lin import array 24*deb4d487SHisping Lin import glob 25*deb4d487SHisping Lin import struct 26*deb4d487SHisping Lin import uuid 27*deb4d487SHisping Lin try: 28*deb4d487SHisping Lin from Cryptodome.Hash import SHA256 29*deb4d487SHisping Lin except ImportError: 30*deb4d487SHisping Lin from Crypto.Hash import SHA256 31*deb4d487SHisping Lin 32*deb4d487SHisping Lin args = get_args(); 33*deb4d487SHisping Lin 34*deb4d487SHisping Lin path = args.tadir + "/*.ta" 35*deb4d487SHisping Lin files = glob.glob(path) 36*deb4d487SHisping Lin for f in files: 37*deb4d487SHisping Lin print(f) 38*deb4d487SHisping Lin 39*deb4d487SHisping Lin fimg = open(args.out, 'wb') 40*deb4d487SHisping Lin magic = 0x524B5441 #RKTA 41*deb4d487SHisping Lin img_ver = 0x00000001 42*deb4d487SHisping Lin ta_num = len(files) 43*deb4d487SHisping Lin reserve = 0 44*deb4d487SHisping Lin header = struct.pack('<IIII', magic, img_ver, ta_num, reserve) 45*deb4d487SHisping Lin fimg.write(header) 46*deb4d487SHisping Lin 47*deb4d487SHisping Lin g_offset = len(header) + (ta_num * 32) 48*deb4d487SHisping Lin for f in files: 49*deb4d487SHisping Lin fta = open(f, 'rb') 50*deb4d487SHisping Lin tadata = fta.read() 51*deb4d487SHisping Lin fta.close 52*deb4d487SHisping Lin ta_name = os.path.basename(f) 53*deb4d487SHisping Lin ta_name_without_extension = os.path.splitext(ta_name)[0] 54*deb4d487SHisping Lin ta_uuid = uuid.UUID(ta_name_without_extension) 55*deb4d487SHisping Lin ta_uuid_hex = bytes.fromhex(ta_uuid.hex) 56*deb4d487SHisping Lin fimg.write(ta_uuid_hex) 57*deb4d487SHisping Lin ta_offset = struct.pack('<I', g_offset) 58*deb4d487SHisping Lin fimg.write(ta_offset) 59*deb4d487SHisping Lin ta_len = struct.pack('<I', len(tadata)) 60*deb4d487SHisping Lin fimg.write(ta_len) 61*deb4d487SHisping Lin ta_ver_num = len(ta_name_without_extension) - 34; 62*deb4d487SHisping Lin ta_ver = struct.pack('<I', ta_ver_num) 63*deb4d487SHisping Lin fimg.write(ta_ver) 64*deb4d487SHisping Lin ta_reserve = struct.pack('<I', reserve) 65*deb4d487SHisping Lin fimg.write(ta_reserve) 66*deb4d487SHisping Lin g_offset += len(tadata) 67*deb4d487SHisping Lin 68*deb4d487SHisping Lin for f in files: 69*deb4d487SHisping Lin fta = open(f, 'rb') 70*deb4d487SHisping Lin tadata = fta.read() 71*deb4d487SHisping Lin fta.close 72*deb4d487SHisping Lin fimg.write(tadata) 73*deb4d487SHisping Lin fimg.close 74*deb4d487SHisping Lin 75*deb4d487SHisping Lin fimg = open(args.out, 'rb+') 76*deb4d487SHisping Lin userta = fimg.read() 77*deb4d487SHisping Lin h = SHA256.new() 78*deb4d487SHisping Lin h.update(userta) 79*deb4d487SHisping Lin img_digest = h.digest() 80*deb4d487SHisping Lin fimg.write(img_digest) 81*deb4d487SHisping Lin fimg.close 82*deb4d487SHisping Lin 83*deb4d487SHisping Lin print("pack ta binary success!") 84*deb4d487SHisping Lin 85*deb4d487SHisping Linif __name__ == "__main__": 86*deb4d487SHisping Lin main() 87