Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 1 | #!/usr/bin/env python3 |
Patrick Georgi | 1afe286 | 2020-05-10 17:34:15 +0200 | [diff] [blame] | 2 | # SPDX-License-Identifier: BSD-2-Clause |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 3 | |
| 4 | import struct |
| 5 | import sys |
| 6 | from io import SEEK_SET, SEEK_END |
| 7 | |
| 8 | class IDBTool: |
| 9 | def __init__(self): |
Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 10 | print("Initialize IDBTool") |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 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 |
Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 28 | buf[x] = struct.pack('B', buf[x] ^ S[k])[0] |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 29 | |
huang lin | 1129f7f | 2016-03-02 16:02:45 +0800 | [diff] [blame] | 30 | def makeIDB(self, chip, from_file, to_file, rc4_flag = False, align_flag = False): |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 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 |
Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 48 | sectors = (data_len + 4 - 1) // SECTOR_SIZE + 1 |
| 49 | pages = (sectors - 1) // PAGE_ALIGN + 1 |
| 50 | sectors = pages * PAGE_ALIGN |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 51 | |
Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 52 | buf = bytearray(sectors * SECTOR_SIZE) |
| 53 | assert len(chip) == 4 |
| 54 | buf[:4] = chip.encode('ascii') |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 55 | buf[4 : 4+data_len] = data |
| 56 | |
Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 57 | idblock = bytearray(4 * SECTOR_SIZE) |
| 58 | blank = bytearray(4 * SECTOR_SIZE) |
| 59 | idblock[:4] = b'\x55\xAA\xF0\x0F' |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 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 | |
Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 69 | idblock[12:16] = struct.pack("<HH", 4, 4) |
| 70 | idblock[506:510] = struct.pack("<HH", sectors, sectors) |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 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): |
Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 80 | fout.write(idblock) |
| 81 | fout.write(blank) |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 82 | |
Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 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) |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 86 | else: |
Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 87 | fout.write(idblock) |
| 88 | fout.write(buf) |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 89 | fout.flush() |
| 90 | except: |
| 91 | sys.exit("Failed to write data to : " + to_file) |
| 92 | finally: |
| 93 | fout.close() |
Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 94 | print("DONE") |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 95 | |
| 96 | def usage(): |
Yilin Yang | 46eaa5a | 2020-09-16 16:01:40 +0800 | [diff] [blame] | 97 | print("Usage: make_idb.py [--chip=RKXX] [--enable-rc4] [--enable-align] [--to=out] --from=in") |
| 98 | print(" --chip: default is RK32") |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 99 | |
| 100 | if __name__ == '__main__': |
| 101 | rc4_flag = align_flag = False |
| 102 | to_file = "IDBlock.bin" |
huang lin | 1129f7f | 2016-03-02 16:02:45 +0800 | [diff] [blame] | 103 | chip = "RK32" |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 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] |
huang lin | 1129f7f | 2016-03-02 16:02:45 +0800 | [diff] [blame] | 114 | elif (para.startswith("--chip=")): |
| 115 | chip = para.split('=')[1] |
huang lin | 817e455 | 2014-08-26 17:31:28 +0800 | [diff] [blame] | 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() |
huang lin | 1129f7f | 2016-03-02 16:02:45 +0800 | [diff] [blame] | 127 | idbtool.makeIDB(chip, from_file, to_file, rc4_flag, align_flag) |