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