1#!/usr/bin/python
2
3from __future__ import print_function
4
5import argparse
6import codecs
7import operator
8import os
9import sys
10from subprocess import Popen, PIPE
11from tempfile import mkstemp
12
13
14def check_sparse(filename):
15    magic = 3978755898
16    with open(filename, 'rb') as i:
17        word = i.read(4)
18        if magic == int(codecs.encode(word[::-1], 'hex'), 16):
19            return True
20    return False
21
22
23def shell_command(comm_list):
24    command = Popen(comm_list, stdout=PIPE, stderr=PIPE)
25    command.communicate()
26    command.wait()
27    if command.returncode != 0:
28        sys.exit(1)
29
30
31def parse_input(input_file):
32    parsed_lines = list()
33    lines = input_file.readlines()
34    for line in lines:
35        line = line.strip()
36        if not line or line[0] == "#":
37            continue
38        params = line.split()
39        if len(params) == 3:
40            for param in params:
41                # interprete file paths such as $OUT/system.img
42                param = os.path.expandvars(param)
43            parsed_lines.append(params)
44
45    partitions = list()
46    num_used = set()
47    for line in parsed_lines:
48        partition_info = dict()
49        partition_info["path"] = line[0]
50        partition_info["label"] = line[1]
51        # round up by 1M
52        sizeByMb = str((1024 * 1024 - 1 + os.path.getsize(line[0])) / 1024 / 1024)
53        partition_info["sizeByMb"] = sizeByMb
54
55        try:
56            partition_info["num"] = int(line[2])
57        except ValueError:
58            print("'%s' cannot be converted to int" % (line[2]))
59            sys.exit(1)
60
61        # check if the partition number is out of range
62        if partition_info["num"] > len(lines) or partition_info["num"] < 0:
63            print("Invalid partition number: %d, range [1..%d]" % \
64                  (partition_info["num"], len(lines)))
65            sys.exit(1)
66
67        # check if the partition number is duplicated
68        if partition_info["num"] in num_used:
69            print("Duplicated partition number:%d" % (partition_info["num"]))
70            sys.exit(1)
71        num_used.add(partition_info["num"])
72        partitions.append(partition_info)
73
74    partitions.sort(key=operator.itemgetter("num"))
75    return partitions
76
77
78def write_partition(partition, output_file, offset):
79    # $ dd if=/path/to/image of=/path/to/output conv=notrunc,sync \
80    #   ibs=1024k obs=1024k seek=<offset>
81    dd_comm = ['dd', 'if=' + partition["path"], 'of=' + output_file, 'conv=notrunc,sync',
82               'ibs=1024k', 'obs=1024k', 'seek=' + str(offset)]
83    shell_command(dd_comm)
84    return
85
86
87def unsparse_partition(partition):
88    # if the input image is in sparse format, unsparse it
89    simg2img = os.environ.get('SIMG2IMG', 'simg2img')
90    print("Unsparsing %s" % (partition["path"]), end=' ')
91    partition["fd"], temp_file = mkstemp()
92    shell_command([simg2img, partition["path"], temp_file])
93    partition["path"] = temp_file
94    print("Done")
95    return
96
97
98def clear_partition_table(filename):
99    sgdisk = os.environ.get('SGDISK', 'sgdisk')
100    print("%s --clear %s" % (sgdisk, filename))
101    shell_command([sgdisk, '--clear', filename])
102    return
103
104
105def add_partition(partition, output_file):
106    sgdisk = os.environ.get('SGDISK', 'sgdisk')
107    num = str(partition["num"])
108    new_comm = '--new=' + num + ':' + partition["start"] + ':' + partition["end"]
109    type_comm = '--type=' + num + ':8300'
110    name_comm = '--change-name=' + num + ':' + partition["label"]
111    # build partition table in order. for example:
112    # $ sgdisk --new=1:2048:5244927 --type=1:8300 --change-name=1:system \
113    #   /path/to/output
114    shell_command([sgdisk, new_comm, type_comm, name_comm, output_file])
115    return
116
117
118def main():
119    # check usage:
120    parser = argparse.ArgumentParser()
121    parser.add_argument("-i", "--input",
122                        type=str, help="input configuration file",
123                        default="image_config")
124    parser.add_argument("-o", "--output",
125                        type=str, help="output filename",
126                        default=os.environ.get("OUT", ".") + "/combined.img")
127    args = parser.parse_args()
128
129    output_filename = os.path.expandvars(args.output)
130
131    # remove the output_filename.qcow2
132    print("removing " + output_filename + ".qcow2")
133    shell_command(['rm', '-rf', output_filename + ".qcow2"])
134
135    # check input file
136    config_filename = args.input
137    if not os.path.exists(config_filename):
138        print("Invalid config file name " + config_filename)
139        sys.exit(1)
140
141    # read input file
142    config = open(config_filename, "r")
143    partitions = parse_input(config)
144    config.close()
145
146    # take a shortcut in build environment
147    if os.path.exists(output_filename) and len(partitions) == 2:
148        print("updating " + output_filename + " ...")
149        shell_command(['dd', "if=" + partitions[0]["path"], "of=" + output_filename,
150                       "conv=notrunc,sync", "ibs=1024k", "obs=1024k", "seek=1"])
151        shell_command(['dd', "if=" + partitions[1]["path"], "of=" + output_filename,
152                       "conv=notrunc,sync", "ibs=1024k", "obs=1024k", "seek=2"])
153        print("done")
154        sys.exit(0)
155    elif len(partitions) == 2:
156        gptprefix = partitions[0]["sizeByMb"] + "_" + partitions[1]["sizeByMb"]
157        prebuilt_gpt_dir = os.path.dirname(os.path.abspath( __file__ )) + "/prebuilt/gpt/" + gptprefix
158        gpt_head = prebuilt_gpt_dir + "/head.img"
159        gpt_tail = prebuilt_gpt_dir + "/head.img"
160        if os.path.exists(gpt_head) and os.path.exists(gpt_tail):
161            print("found prebuilt gpt header and footer, use it")
162            shell_command(['dd', "if=" + gpt_head, "of=" + output_filename, "bs=1024k",
163                    "conv=notrunc,sync", "count=1"])
164            shell_command(['dd', "if=" + partitions[0]["path"], "of=" + output_filename,
165                    "bs=1024k", "conv=notrunc,sync", "seek=1"])
166            shell_command(['dd', "if=" + partitions[1]["path"], "of=" + output_filename,
167                    "bs=1024k", "conv=notrunc,sync", "seek=" + str(1 + int(partitions[0]["sizeByMb"]))])
168            shell_command(['dd', "if=" + gpt_tail, "of=" + output_filename,
169                    "bs=1024k", "conv=notrunc,sync",
170                    "seek=" + str(1 + int(partitions[0]["sizeByMb"]) + int(partitions[1]["sizeByMb"]))])
171            print("done")
172            sys.exit(0)
173
174    # combine the images
175    # add padding
176    shell_command(['dd', 'if=/dev/zero', 'of=' + output_filename, 'ibs=1024k', 'count=1'])
177
178    for partition in partitions:
179        offset = os.path.getsize(output_filename)
180        partition["start"] = str(offset // 512)
181        # dectect sparse file format
182        if check_sparse(partition["path"]):
183            unsparse_partition(partition)
184
185        # TODO: extract the partition if the image file is already formatted
186
187        write_partition(partition, output_filename, offset // 1024 // 1024)
188        offset = os.path.getsize(output_filename)
189        partition["end"] = str(offset // 512 - 1)
190
191    # add padding
192    # $ dd if=/dev/zero of=/path/to/output conv=notrunc bs=1 \
193    #   count=1024k seek=<offset>
194    offset = os.path.getsize(output_filename) // 1024 // 1024
195    shell_command(['dd', 'if=/dev/zero', 'of=' + output_filename,
196                   'conv=notrunc', 'bs=1024k', 'count=1', 'seek=' + str(offset)])
197
198    # make partition table
199    # $ sgdisk --clear /path/to/output
200    clear_partition_table(output_filename)
201
202    for partition in partitions:
203        add_partition(partition, output_filename)
204        # clean up, delete any unsparsed image files generated
205        if 'fd' in partition:
206            os.close(partition["fd"])
207            os.remove(partition["path"])
208
209
210if __name__ == "__main__":
211    main()
212