1#!/usr/bin/env python
2
3import os
4import os.path
5import sys, getopt
6import binascii
7import struct
8import string
9
10class generator(object):
11    #
12    # struct l_loader_head {
13    #      unsigned int	first_instr;
14    #      unsigned char	magic[16];	@ BOOTMAGICNUMBER!
15    #      unsigned int	l_loader_start;
16    #      unsigned int	l_loader_end;
17    # };
18    file_header = [0, 0, 0, 0, 0, 0, 0]
19
20    #
21    # struct entry_head {
22    #       unsigned char   magic[8];           @ ENTY
23    #       unsigned char   name[8];            @ loader/bl1
24    #       unsigned int    start_lba;
25    #       unsigned int    count_lba;
26    #       unsigned int    flag;               @ boot partition or not
27    # };
28
29    s1_entry_name = ['loader', 'bl1']
30    s2_entry_name = ['primary', 'second']
31
32    block_size = 512
33
34    stage = 0
35
36    # set in self.add()
37    idx = 0
38
39    # set in self.parse()
40    ptable_lba = 0
41    stable_lba = 0
42
43    # file pointer
44    p_entry = 0
45    p_file = 0
46
47    def __init__(self, out_img):
48        try:
49            self.fp = open(out_img, "wb+")
50        except IOError, e:
51            print "*** file open error:", e
52            sys.exit(3)
53        else:
54            self.entry_hd = [[0 for col in range(7)] for row in range(5)]
55
56    def __del__(self):
57        self.fp.close()
58
59    # parse partition from the primary ptable
60    def parse(self, fname):
61        try:
62            fptable = open(fname, "rb")
63        except IOError, e:
64            print "*** file open error:", e
65            sys.exit(3)
66        else:
67            # skip the first block in primary partition table
68            # that is MBR protection information
69            fptable.read(self.block_size)
70            # check whether it's a primary paritition table
71            data = struct.unpack("8s", fptable.read(8))
72            efi_magic = 'EFI PART'
73            if cmp("EFI PART", data[0]):
74                print "It's not partition table image."
75                fptable.close()
76                sys.exit(4)
77            # skip 16 bytes
78            fptable.read(16)
79            # get lba of both primary partition table and secondary partition table
80            data = struct.unpack("QQQQ", fptable.read(32))
81            self.ptable_lba = data[0] - 1
82            self.stable_lba = data[3] + 1
83            # skip 24 bytes
84            fptable.read(24)
85            data = struct.unpack("i", fptable.read(4))
86            pentries = data[0]
87            # skip the reset in this block
88            fptable.read(self.block_size - 84)
89
90            for i in range(1, pentries):
91                # name is encoded as UTF-16
92                d0,lba,d2,name = struct.unpack("32sQ16s72s", fptable.read(128))
93                plainname = unicode(name, "utf-16")
94                if (not cmp(plainname[0:7], 'l-loader'[0:7])):
95                    print 'bl1_lba: ', lba
96                    self.bl1_lba = lba
97                    sys.exit(1)
98
99            fptable.close()
100
101    def add(self, lba, fname):
102        try:
103            fsize = os.path.getsize(fname)
104        except IOError, e:
105            print "*** file open error:", e
106            sys.exit(4)
107        else:
108            blocks = (fsize + self.block_size - 1) / self.block_size
109            if (self.stage == 1):
110                # Boot Area1 in eMMC
111                bootp = 1
112                if self.idx == 0:
113                    self.p_entry = 28
114            elif (self.stage == 2):
115                # User Data Area in eMMC
116                bootp = 0
117                # create an empty block only for stage2
118                # This empty block is used to store entry head
119                print 'p_file: ', self.p_file, 'p_entry: ', self.p_entry
120                if self.idx == 0:
121                    self.fp.seek(self.p_file)
122                    for i in range (0, self.block_size):
123                        zero = struct.pack('x')
124                        self.fp.write(zero)
125                    self.p_file += self.block_size
126                    self.p_entry = 0
127            else:
128                print "wrong stage ", stage, "is specified"
129                sys.exit(4)
130            # Maybe the file size isn't aligned. So pad it.
131            if (self.idx == 0) and (self.stage == 1):
132                if fsize > 2048:
133                    print 'loader size exceeds 2KB. file size: ', fsize
134                    sys.exit(4)
135                else:
136                    left_bytes = 2048 - fsize
137            else:
138                left_bytes = fsize % self.block_size
139                if left_bytes:
140                    left_bytes = self.block_size - left_bytes
141            print 'lba: ', lba, 'blocks: ', blocks, 'bootp: ', bootp, 'fname: ', fname
142            # write images
143            fimg = open(fname, "rb")
144            for i in range (0, blocks):
145                buf = fimg.read(self.block_size)
146                self.fp.seek(self.p_file)
147                self.fp.write(buf)
148                # p_file is the file pointer of the new binary file
149                # At last, it means the total block size of the new binary file
150                self.p_file += self.block_size
151
152            if (self.idx == 0) and (self.stage == 1):
153                self.p_file = 2048
154            print 'p_file: ', self.p_file, 'last block is ', fsize % self.block_size, 'bytes', '  tell: ', self.fp.tell(), 'left_bytes: ', left_bytes
155            if left_bytes:
156                for i in range (0, left_bytes):
157                    zero = struct.pack('x')
158                    self.fp.write(zero)
159                print 'p_file: ', self.p_file, '  pad to: ', self.fp.tell()
160
161            # write entry information at the header
162            if self.stage == 1:
163                byte = struct.pack('8s8siii', 'ENTRYHDR', self.s1_entry_name[self.idx], lba, blocks, bootp)
164            elif self.stage == 2:
165                byte = struct.pack('8s8siii', 'ENTRYHDR', self.s2_entry_name[self.idx], lba, blocks, bootp)
166            self.fp.seek(self.p_entry)
167            self.fp.write(byte)
168            self.p_entry += 28
169            self.idx += 1
170
171            fimg.close()
172
173    def hex2(self, data):
174        return data > 0 and hex(data) or hex(data & 0xffffffff)
175
176    def end(self):
177        if self.stage == 1:
178            self.fp.seek(20)
179            start,end = struct.unpack("ii", self.fp.read(8))
180            print "start: ", self.hex2(start), 'end: ', self.hex2(end)
181            end = start + self.p_file
182            print "start: ", self.hex2(start), 'end: ', self.hex2(end)
183            self.fp.seek(24)
184            byte = struct.pack('i', end)
185            self.fp.write(byte)
186        self.fp.close()
187
188    def create_stage1(self, img_loader, img_bl1, output_img):
189        print '+-----------------------------------------------------------+'
190        print ' Input Images:'
191        print '     loader:                       ', img_loader
192        print '     bl1:                          ', img_bl1
193        print ' Ouput Image:                      ', output_img
194        print '+-----------------------------------------------------------+\n'
195
196        self.stage = 1
197
198        # The first 2KB is reserved
199        # The next 2KB is for loader image
200        self.add(4, img_loader)    # img_loader doesn't exist in partition table
201        print 'self.idx: ', self.idx
202        # bl1.bin starts from 4KB
203        self.add(8, img_bl1)      # img_bl1 doesn't exist in partition table
204
205    def create_stage2(self, img_prm_ptable, img_sec_ptable, output_img):
206        print '+-----------------------------------------------------------+'
207        print ' Input Images:'
208        print '     primary partition table:      ', img_prm_ptable
209        print '     secondary partition table:    ', img_sec_ptable
210        print ' Ouput Image:                      ', output_img
211        print '+-----------------------------------------------------------+\n'
212
213        self.stage = 2
214        self.parse(img_prm_ptable)
215        self.add(self.ptable_lba, img_prm_ptable)
216        if (cmp(img_sec_ptable, 'secondary partition table')):
217            # Doesn't match. It means that secondary ptable is specified.
218            self.add(self.stable_lba, img_sec_ptable)
219        else:
220            print 'Don\'t need secondary partition table'
221
222def main(argv):
223    stage1 = 0
224    stage2 = 0
225    img_prm_ptable = "primary partition table"
226    img_sec_ptable = "secondary partition table"
227    try:
228        opts, args = getopt.getopt(argv,"ho:",["img_loader=","img_bl1=","img_prm_ptable=","img_sec_ptable="])
229    except getopt.GetoptError:
230        print 'gen_loader.py -o <l-loader.bin> --img_loader <l-loader> --img_bl1 <bl1.bin> --img_prm_ptable <prm_ptable.img> --img_sec_ptable <sec_ptable.img>'
231        sys.exit(2)
232    for opt, arg in opts:
233        if opt == '-h':
234            print 'gen_loader.py -o <l-loader.bin> --img_loader <l-loader> --img_bl1 <bl1.bin> --img_prm_ptable <prm_ptable.img> --img_sec_ptable <sec_ptable.img>'
235            sys.exit(1)
236        elif opt == '-o':
237            output_img = arg
238        elif opt in ("--img_loader"):
239            img_loader = arg
240            stage1 = 1
241        elif opt in ("--img_bl1"):
242            img_bl1 = arg
243            stage1 = 1
244        elif opt in ("--img_prm_ptable"):
245            img_prm_ptable = arg
246            stage2 = 1
247        elif opt in ("--img_sec_ptable"):
248            img_sec_ptable = arg
249
250    loader = generator(output_img)
251    loader.idx = 0
252
253    if (stage1 == 1) and (stage2 == 1):
254        print 'There are only loader & BL1 in stage1.'
255        print 'And there are primary partition table, secondary partition table and FIP in stage2.'
256        sys.exit(1)
257    elif (stage1 == 0) and (stage2 == 0):
258        print 'No input images are specified.'
259        sys.exit(1)
260    elif stage1 == 1:
261        loader.create_stage1(img_loader, img_bl1, output_img)
262    elif stage2 == 1:
263        loader.create_stage2(img_prm_ptable, img_sec_ptable, output_img)
264
265    loader.end()
266
267if __name__ == "__main__":
268    main(sys.argv[1:])
269