1#!/usr/bin/env python
2# coding: utf-8
3# Copyright (c) 2022 Huawei Device Co., Ltd.
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.
15
16import sys
17import os
18import hashlib
19import errno
20import stat
21import datetime
22
23
24def usage():
25    print('\n Usage: imgcovert.py <cmd> <input> <output>')
26    print('         <cmd>: sparse or unsparse')
27    print('         <input>: input image file')
28    print('         <output>: ouput image file\n')
29    return
30
31
32def get_fill_cnt(inputfile: str, blocksize: int) -> int:
33    flags = os.O_WRONLY
34    modes = stat.S_IWUSR | stat.S_IRUSR
35    size = os.path.getsize(inputfile)
36    fill_cnt = 0
37    if size % blocksize != 0:
38        fill_cnt = blocksize - size % blocksize
39    indata = os.fdopen(os.open(inputfile, flags, modes), 'a')
40    for _ in range(fill_cnt):
41        indata.write("\0")
42    indata.close()
43    return fill_cnt
44
45
46def get_gap_blocksize(length: int, size: int) -> int:
47    if length < size:
48        cnt = 2
49    elif length < (size * 2):
50        cnt = 3
51    else:
52        cnt = 4
53    return cnt
54
55
56def get_block_cnt(inputfile: str, blocksize: int) -> int:
57    size = os.path.getsize(inputfile)
58    if blocksize != 0:
59        totalblocks = size / blocksize
60    else:
61        sys.exit(1)
62    if (size % blocksize) != 0:
63        print("len is not eq n * blocksize: ", size, totalblocks)
64    return totalblocks
65
66
67def get_crc_value(inputfile: str, blocksize: int):
68    totalblocks = get_block_cnt(inputfile, blocksize)
69    with open(inputfile, 'rb') as indata:
70        ind = 0
71        md5 = hashlib.md5()
72        while (ind < totalblocks):
73            md5.update(indata.read(blocksize))
74            ind += 1
75    return md5.hexdigest()
76
77
78def unsparse(sparseimagefile: str, imagefile: str):
79    with open(sparseimagefile, 'r') as header:
80        magic_mumber = header.readline()
81        version = header.readline()
82        blocksize = int(header.readline())
83        total_blocks = int(header.readline())
84        crc_value = header.readline()
85        input_crc_value = header.readline()
86        table_numbers = int(header.readline())
87        table = []
88        flags = os.O_CREAT | os.O_RDWR
89        modes = stat.S_IWUSR | stat.S_IRUSR
90        i = 0
91        while (i < table_numbers):
92            start = int(header.readline())
93            end = int(header.readline())
94            table.append([start, end])
95            i += 1
96        fill_cnt = int(header.readline())
97        length = header.tell()
98    with open(sparseimagefile, 'rb') as inputrow:
99        inputrow.seek(get_gap_blocksize(length, blocksize) * blocksize)
100        output = os.fdopen(os.open(imagefile, flags, modes), 'wb')
101        output.truncate(total_blocks * blocksize)
102        md5 = hashlib.md5()
103        for block in table:
104            cnt = block[1] - block[0]
105            output.seek(block[0] * blocksize)
106            indata = inputrow.read(cnt * blocksize)
107            md5.update(indata)
108            output.write(indata)
109        output.close()
110    print("RawFileCRC: ", get_crc_value(imagefile, blocksize), crc_value)
111    print("SparseCRC: ", md5.hexdigest(), input_crc_value)
112    with open(imagefile, 'r+') as output:
113        output.truncate(total_blocks * blocksize - fill_cnt)
114    return
115
116
117def is_empty_block(buff: list, size: int) -> bool:
118    ind = 0
119    while (ind < size):
120        if buff[ind] != 0:
121            return False
122        ind += 1
123    return True
124
125
126def process_block(inputrow, blocksize, outputtemp, blockid, total_blocks) -> int:
127    ind = 0
128    start = -1
129    table_numbers = 0
130    while (ind < total_blocks):
131        indata = inputrow.read(blocksize)
132        if len(indata) != blocksize:
133            print("error Block", ind, len(indata))
134        if is_empty_block(indata, blocksize):
135            if start != -1:
136                blockid.append([start, ind])
137                table_numbers += 1
138                start = -1
139        else:
140            outputtemp.write(indata)
141            if start == -1:
142                start = ind
143        ind += 1
144    if start != -1:
145        blockid.append([start, ind])
146        table_numbers += 1
147        start = -1
148    return table_numbers
149
150
151def get_raw_datafile(imagefile: str, blockid, total_blocks: int, blocksize: int) -> int:
152    temp_file = imagefile + ".tempfile"
153    flags = os.O_CREAT | os.O_RDWR
154    modes = stat.S_IWUSR | stat.S_IRUSR
155    table_numbers = 0
156    with open(imagefile, 'rb') as inputrow, os.fdopen(os.open(temp_file, flags, modes), 'wb') as outputtemp:
157        table_numbers = process_block(inputrow, blocksize, outputtemp, blockid, total_blocks)
158    return table_numbers
159
160
161def sparse(imagefile: str, sparseimagefile: str):
162    temp_file = imagefile + ".tempfile"
163    magic_number = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
164    version = 1.0
165    blocksize = 4096
166    table_numbers = 0
167    blockid = []
168    flags = os.O_CREAT | os.O_RDWR
169    modes = stat.S_IWUSR | stat.S_IRUSR
170
171    fill_cnt = get_fill_cnt(imagefile, blocksize)
172    total_blocks = get_block_cnt(imagefile, blocksize)
173    table_numbers = get_raw_datafile(imagefile, blockid, total_blocks, blocksize)
174
175#   save the header
176    outputrow = os.fdopen(os.open(sparseimagefile, flags, modes), 'w')
177    outputrow.write("%s\n" % (magic_number))
178    outputrow.write("%s\n" % (version))
179    outputrow.write("%s\n" % (blocksize))
180    outputrow.write("%s\n" % (int(total_blocks)))
181    outputrow.write("%s\n" % (get_crc_value(imagefile, blocksize)))
182    outputrow.write("%s\n" % (get_crc_value(temp_file, blocksize)))
183    outputrow.write("%s\n" % (table_numbers))
184    for block in blockid:
185        outputrow.write("%s\n" % (block[0]))
186        outputrow.write("%s\n" % (block[1]))
187    outputrow.write("%s\n" % (int(fill_cnt)))
188    outputrow.truncate(get_gap_blocksize(outputrow.tell(), blocksize) * blocksize)
189    outputrow.close()
190
191#   append the raw data
192    outputrow = os.fdopen(os.open(sparseimagefile, flags, modes), 'ab')
193    outputtemp = os.fdopen(os.open(temp_file, flags, modes), 'rb')
194    blocknum = get_block_cnt(temp_file, blocksize)
195    i = 0
196    while (i < blocknum):
197        outputrow.write(outputtemp.read(blocksize))
198        i += 1
199    outputtemp.close()
200    outputrow.close()
201    os.remove(temp_file)
202    with open(imagefile, 'r+') as output:
203        output.truncate(int(total_blocks) * int(blocksize) - int(fill_cnt))
204
205
206if __name__ == '__main__':
207    if len(sys.argv) != 4:
208        usage()
209        sys.exit()
210    CMD = str(sys.argv[1])
211    INPUT_FILE = str(sys.argv[2])
212    OUTPUT_FILE = str(sys.argv[3])
213    if CMD == 'unsparse':
214        unsparse(INPUT_FILE, OUTPUT_FILE)
215    elif CMD == 'sparse':
216        sparse(INPUT_FILE, OUTPUT_FILE)
217    else:
218        usage()