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