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