1#!/usr/bin/env python
2
3import sys
4import os
5
6import struct
7from collections import namedtuple
8import io
9
10import argparse
11parser = argparse.ArgumentParser()
12parser.add_argument('rootdir', type=str, help='the path to rootfs')
13parser.add_argument('output', type=argparse.FileType('wb'), nargs='?', help='output file name')
14parser.add_argument('--dump', action='store_true', help='dump the fs hierarchy')
15parser.add_argument('--binary', action='store_true', help='output binary file')
16parser.add_argument('--addr', default='0', help='set the base address of the binary file, default to 0.')
17
18class File(object):
19    def __init__(self, name):
20        self._name = name
21        self._data = open(name, 'rb').read()
22
23    @property
24    def name(self):
25        return self._name
26
27    @property
28    def c_name(self):
29        return '_' + self._name.replace('.', '_')
30
31    @property
32    def bin_name(self):
33        # Pad to 4 bytes boundary with \0
34        pad_len = 4
35        bn = self._name + '\0' * (pad_len - len(self._name) % pad_len)
36        return bn
37
38    def c_data(self, prefix=''):
39        '''Get the C code represent of the file content.'''
40        head = 'static const rt_uint8_t %s[] = {\n' % \
41                (prefix + self.c_name)
42        tail = '\n};'
43
44        if self.entry_size == 0:
45            return ''
46        if len(self._data) > 0 and type(self._data[0]) == int:
47            return head + ','.join(('0x%02x' % i for i in self._data)) + tail
48        else:
49            return head + ','.join(('0x%02x' % ord(i) for i in self._data)) + tail
50
51    @property
52    def entry_size(self):
53        return len(self._data)
54
55    def bin_data(self, base_addr=0x0):
56        return bytes(self._data)
57
58    def dump(self, indent=0):
59        print('%s%s' % (' ' * indent, self._name))
60
61class Folder(object):
62    bin_fmt = struct.Struct('IIII')
63    bin_item = namedtuple('dirent', 'type, name, data, size')
64
65    def __init__(self, name):
66        self._name = name
67        self._children = []
68
69    @property
70    def name(self):
71        return self._name
72
73    @property
74    def c_name(self):
75        # add _ to avoid conflict with C key words.
76        return '_' + self._name
77
78    @property
79    def bin_name(self):
80        # Pad to 4 bytes boundary with \0
81        pad_len = 4
82        bn = self._name + '\0' * (pad_len - len(self._name) % pad_len)
83        return bn
84
85    def walk(self):
86        # os.listdir will return unicode list if the argument is unicode.
87        # TODO: take care of the unicode names
88        for ent in os.listdir(u'.'):
89            if os.path.isdir(ent):
90                cwd = os.getcwd()
91                d = Folder(ent)
92                # depth-first
93                os.chdir(os.path.join(cwd, ent))
94                d.walk()
95                # restore the cwd
96                os.chdir(cwd)
97                self._children.append(d)
98            else:
99                self._children.append(File(ent))
100
101    def sort(self):
102        def _sort(x, y):
103            if x.name == y.name:
104                return 0
105            elif x.name > y.name:
106                return 1
107            else:
108                return -1
109        from functools import cmp_to_key
110        self._children.sort(key=cmp_to_key(_sort))
111
112        # sort recursively
113        for c in self._children:
114            if isinstance(c, Folder):
115                c.sort()
116
117    def dump(self, indent=0):
118        print('%s%s' % (' ' * indent, self._name))
119        for c in self._children:
120            c.dump(indent + 1)
121
122    def c_data(self, prefix=''):
123        '''get the C code represent of the folder.
124
125           It is recursive.'''
126        # make the current dirent
127        # static is good. Only root dirent is global visible.
128        if self.entry_size == 0:
129            return ''
130
131        dhead = 'static const struct romfs_dirent %s[] = {\n' % (prefix + self.c_name)
132        dtail = '\n};'
133        body_fmt = '    {{{type}, "{name}", (rt_uint8_t *){data}, sizeof({data})/sizeof({data}[0])}}'
134        body_fmt0= '    {{{type}, "{name}", RT_NULL, 0}}'
135        # prefix of children
136        cpf = prefix+self.c_name
137        body_li = []
138        payload_li = []
139        for c in self._children:
140            entry_size = c.entry_size
141            if isinstance(c, File):
142                tp = 'ROMFS_DIRENT_FILE'
143            elif isinstance(c, Folder):
144                tp = 'ROMFS_DIRENT_DIR'
145            else:
146                assert False, 'Unkown instance:%s' % str(c)
147            if entry_size == 0:
148                body_li.append(body_fmt0.format(type=tp, name = c.name))
149            else:
150                body_li.append(body_fmt.format(type=tp,
151                                            name=c.name,
152                                            data=cpf+c.c_name))
153            payload_li.append(c.c_data(prefix=cpf))
154
155        # All the data we need is defined in payload so we should append the
156        # dirent to it. It also meet the depth-first policy in this code.
157        payload_li.append(dhead + ',\n'.join(body_li) + dtail)
158
159        return '\n\n'.join(payload_li)
160
161    @property
162    def entry_size(self):
163        return len(self._children)
164
165    def bin_data(self, base_addr=0x0):
166        '''Return StringIO object'''
167        # The binary layout is different from the C code layout. We put the
168        # dirent before the payload in this mode. But the idea is still simple:
169        #                           Depth-First.
170
171        #{
172        #  rt_uint32_t type;
173        #  const char *name;
174        #  const rt_uint8_t *data;
175        #  rt_size_t size;
176        #}
177        d_li = []
178        # payload base
179        p_base = base_addr + self.bin_fmt.size * self.entry_size
180        # the length to record how many data is in
181        v_len = p_base
182        # payload
183        p_li = []
184        for c in self._children:
185            if isinstance(c, File):
186                # ROMFS_DIRENT_FILE
187                tp = 0
188            elif isinstance(c, Folder):
189                # ROMFS_DIRENT_DIR
190                tp = 1
191            else:
192                assert False, 'Unkown instance:%s' % str(c)
193
194            name = bytes(c.bin_name.encode('utf-8'))
195            name_addr = v_len
196            v_len += len(name)
197
198            data = c.bin_data(base_addr=v_len)
199            data_addr = v_len
200            # pad the data to 4 bytes boundary
201            pad_len = 4
202            if len(data) % pad_len != 0:
203                data += ('\0' * (pad_len - len(data) % pad_len)).encode('utf-8')
204            v_len += len(data)
205
206            d_li.append(self.bin_fmt.pack(*self.bin_item(
207                                               type=tp,
208                                               name=name_addr,
209                                               data=data_addr,
210                                               size=c.entry_size)))
211
212            p_li.extend((name, data))
213
214        return bytes().join(d_li) + bytes().join(p_li)
215
216def get_c_data(tree):
217    # Handle the root dirent specially.
218    root_dirent_fmt = '''/* Generated by mkromfs. Edit with caution. */
219#include <rtthread.h>
220#include <dfs_romfs.h>
221
222{data}
223
224const struct romfs_dirent {name} = {{
225    ROMFS_DIRENT_DIR, "/", (rt_uint8_t *){rootdirent}, sizeof({rootdirent})/sizeof({rootdirent}[0])
226}};
227'''
228
229    return root_dirent_fmt.format(name='romfs_root',
230                                  rootdirent=tree.c_name,
231                                  data=tree.c_data())
232
233def get_bin_data(tree, base_addr):
234    v_len = base_addr + Folder.bin_fmt.size
235    name = bytes('/\0\0\0'.encode("utf-8"))
236    name_addr = v_len
237    v_len += len(name)
238    data_addr = v_len
239    # root entry
240    data = Folder.bin_fmt.pack(*Folder.bin_item(type=1,
241                                                name=name_addr,
242                                                data=data_addr,
243                                                size=tree.entry_size))
244    return data + name + tree.bin_data(v_len)
245
246if __name__ == '__main__':
247    args = parser.parse_args()
248
249    os.chdir(args.rootdir)
250
251    tree = Folder('romfs_root')
252    tree.walk()
253    tree.sort()
254
255    if args.dump:
256        tree.dump()
257
258    if args.binary:
259        data = get_bin_data(tree, int(args.addr, 16))
260    else:
261        data = get_c_data(tree).encode()
262
263    output = args.output
264    if not output:
265        output = sys.stdout
266
267    output.write(data)
268