1#!/usr/bin/env python3 2# SPDX-License-Identifier: BSD-2-Clause 3 4import struct 5import sys 6from io import SEEK_SET, SEEK_END 7 8class IDBTool: 9 def __init__(self): 10 print("Initialize IDBTool") 11 12 def p_rc4(self, buf, length): 13 key = (124,78,3,4,85,5,9,7,45,44,123,56,23,13,23,17) 14 K = key * 16 15 S = [i for i in range(256)] 16 17 j = 0 18 for i in range(256): 19 j = (j + S[i] + K[i]) % 256 20 temp = S[i]; S[i] = S[j]; S[j] = temp; 21 22 i = j = k = 0 23 for x in range(length): 24 i = (i+1) % 256 25 j = (j + S[i]) % 256 26 temp = S[i]; S[i] = S[j]; S[j] = temp 27 k = (S[i] + S[j]) % 256 28 buf[x] = struct.pack('B', buf[x] ^ S[k])[0] 29 30 def makeIDB(self, chip, from_file, to_file, rc4_flag = False, align_flag = False): 31 try: 32 fin = open(from_file, 'rb') 33 except: 34 sys.exit("Failed to open file : " + from_file) 35 36 try: 37 fin.seek(0, SEEK_END) 38 if (fin.tell() > 4 * 1024 * 1024): 39 sys.exit("Input file is more than 4MB") 40 fin.seek(0) 41 data = fin.read() 42 finally: 43 fin.close() 44 45 data_len = len(data) 46 SECTOR_SIZE = 512 47 PAGE_ALIGN = 4 48 sectors = (data_len + 4 - 1) // SECTOR_SIZE + 1 49 pages = (sectors - 1) // PAGE_ALIGN + 1 50 sectors = pages * PAGE_ALIGN 51 52 buf = bytearray(sectors * SECTOR_SIZE) 53 assert len(chip) == 4 54 buf[:4] = chip.encode('ascii') 55 buf[4 : 4+data_len] = data 56 57 idblock = bytearray(4 * SECTOR_SIZE) 58 blank = bytearray(4 * SECTOR_SIZE) 59 idblock[:4] = b'\x55\xAA\xF0\x0F' 60 61 if (not rc4_flag): 62 idblock[8:12] = struct.pack("<I", 1) 63 else: 64 for i in range(sectors): 65 list_tmp = buf[SECTOR_SIZE*i : SECTOR_SIZE*(i+1)] 66 self.p_rc4(list_tmp, SECTOR_SIZE) 67 buf[SECTOR_SIZE*i : SECTOR_SIZE*(i+1)] = list_tmp 68 69 idblock[12:16] = struct.pack("<HH", 4, 4) 70 idblock[506:510] = struct.pack("<HH", sectors, sectors) 71 self.p_rc4(idblock, SECTOR_SIZE) 72 73 try: 74 fout = open(to_file, "wb+") 75 except: 76 sys.exit("Failed to open output file : " + to_file) 77 78 try: 79 if (align_flag): 80 fout.write(idblock) 81 fout.write(blank) 82 83 for s in range(0, sectors * SECTOR_SIZE, PAGE_ALIGN * SECTOR_SIZE): 84 fout.write(buf[s : s + PAGE_ALIGN * SECTOR_SIZE]) 85 fout.write(blank) 86 else: 87 fout.write(idblock) 88 fout.write(buf) 89 fout.flush() 90 except: 91 sys.exit("Failed to write data to : " + to_file) 92 finally: 93 fout.close() 94 print("DONE") 95 96def usage(): 97 print("Usage: make_idb.py [--chip=RKXX] [--enable-rc4] [--enable-align] [--to=out] --from=in") 98 print(" --chip: default is RK32") 99 100if __name__ == '__main__': 101 rc4_flag = align_flag = False 102 to_file = "IDBlock.bin" 103 chip = "RK32" 104 105 for para in sys.argv[1:]: 106 if (para == "--enable-rc4"): 107 rc4_flag = True 108 elif (para == "--enable-align"): 109 align_flag = True 110 elif (para.startswith("--to=")): 111 to_file = para.split('=')[1] 112 elif (para.startswith("--from=")): 113 from_file = para.split('=')[1] 114 elif (para.startswith("--chip=")): 115 chip = para.split('=')[1] 116 elif (para == "--help" or para == "-h"): 117 usage() 118 sys.exit() 119 else: 120 usage() 121 sys.exit() 122 if ('from_file' not in vars() or to_file == ''): 123 usage() 124 sys.exit() 125 126 idbtool = IDBTool() 127 idbtool.makeIDB(chip, from_file, to_file, rc4_flag, align_flag) 128