1#!/usr/bin/env python
2# Copyright 2017, The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15from __future__ import print_function
16try:
17	from os import fstat, stat, remove
18	from sys import exit
19	from argparse import ArgumentParser, FileType
20	from ctypes import sizeof, Structure, c_char, c_int
21	from struct import pack, calcsize
22	import zlib
23except Exception as e:
24	print("some module is needed:" + str(e))
25	exit(-1)
26
27dt_head_info_fmt = '4sII'
28dt_entry_fmt = 'Q4I2Q'
29dtimg_version = 1
30dtb_count = 1
31
32def write32(output, value):
33	output.write(chr(value & 255)) ; value=value // 256
34	output.write(chr(value & 255)) ; value=value // 256
35	output.write(chr(value & 255)) ; value=value // 256
36	output.write(chr(value & 255))
37
38def compress(filename, input, output):
39	output.write('\037\213\010')
40	output.write(chr(0))
41
42	statval = stat(filename)
43	write32(output, 0)
44	output.write('\002')
45	output.write('\003')
46
47	crcval = zlib.crc32("")
48	compobj = zlib.compressobj(9, zlib.DEFLATED, -zlib.MAX_WBITS,
49		zlib.DEF_MEM_LEVEL, 0)
50	while True:
51		data = input.read(1024)
52		if data == "":
53			break
54		crcval = zlib.crc32(data, crcval)
55		output.write(compobj.compress(data))
56	output.write(compobj.flush())
57	write32(output, crcval)
58	write32(output, statval.st_size)
59
60def dtb_compress(dtb_file):
61	try:
62		outputname = dtb_file + '.gz'
63		input = open(dtb_file, 'rb')
64		output = open(outputname, 'wb')
65		compress(dtb_file, input, output)
66		input.close()
67		output.close()
68	except Exception as e:
69		print('dtb_compress error:' + str(e))
70		exit(-1)
71	return outputname
72
73class dt_head_info(Structure):
74	_fields_ = [('magic', c_char * 4),
75		    ('version', c_int),
76		    ('dt_count', c_int)]
77
78class dt_entry_t(Structure):
79	_fields_ = [('dtb_size', c_int),
80		    ('dtb_offset', c_int)]
81
82def align_page_size(offset, pagesize):
83	return (pagesize - (offset % pagesize))
84
85def write_head_info(head_info, args):
86	args.output.write(pack(dt_head_info_fmt,
87			       head_info.magic,
88			       head_info.version,
89			       head_info.dt_count))
90
91def write_dtb_entry_t(dt_entry, args):
92	args.output.write(pack(dt_entry_fmt,
93			       0,  # reserved
94			       dt_entry.dtb_size,
95			       0,  # reserved
96			       dt_entry.dtb_offset,
97			       0,  # reserved
98			       0,  # reserved
99			       0)) # reserved
100
101def write_padding(args, padding):
102	for i in range(0, padding):
103		args.output.write('\x00')
104
105def write_dtb(args):
106	dtb_file = args.dtb
107	out_dtb = dtb_file
108	if args.compress == True:
109		out_dtb = dtb_compress(dtb_file)
110	try:
111		dtb_offset = calcsize(dt_head_info_fmt) + \
112				      calcsize(dt_entry_fmt) + \
113				      4
114		padding = align_page_size(dtb_offset, args.pagesize)
115		dtb_size = stat(out_dtb).st_size
116		dtb_size_padding = align_page_size(dtb_size, args.pagesize)
117		dt_entry = dt_entry_t(dtb_size + dtb_size_padding,
118				      dtb_offset + padding)
119		write_dtb_entry_t(dt_entry, args)
120		args.output.write(pack('I', 0)) # SUCCESS code number
121		write_padding(args, padding)
122		with open(out_dtb, 'rb') as dtb_fd:
123			args.output.write(dtb_fd.read(dtb_size))
124			write_padding(args, dtb_size_padding)
125	except Exception as e:
126		print('write dtb error:' + str(e))
127		exit(-1)
128
129def clean_gz_file(args):
130	try:
131		if args.compress != True:
132			return
133		remove(args.dtb + '.gz')
134	except Exception as e:
135		print('clean gz file error:' + str(e))
136		exit(-1)
137
138def parse_cmdline():
139	parser = ArgumentParser()
140	parser.add_argument('-c', '--compress', help='compress dtb or not',
141			    action='store_true')
142	parser.add_argument('-d', '--dtb', help='path to the dtb', type=str,
143			    required=True)
144	parser.add_argument('-s', '--pagesize', help='align page size',
145			    type=int, choices=[2**i for i in range(11,15)],
146			    default=2048)
147	parser.add_argument('-o', '--output', help='output file name',
148			    type=FileType('wb'), required=True)
149	return parser.parse_args()
150
151def main():
152	args = parse_cmdline()
153	dtimg_head_info = dt_head_info('HSDT', dtimg_version, dtb_count)
154	write_head_info(dtimg_head_info, args)
155	write_dtb(args)
156	clean_gz_file(args)
157
158if __name__ == '__main__':
159	main()
160