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