awesome-android/avbroot
3
0
Fork 0
mirror of https://github.com/chenxiaolong/avbroot.git synced 2026-07-03 14:05:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

Replace double quotes with single quotes where possible

Browse Source 
Also formatted differently to adhere more to the style used in the rest
of the project
This commit is contained in:
Pascal Roeleven
2023-02-28 20:56:46 +01:00
parent 13b9877e2c
commit a7170ea4c0
2 changed files with 137 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests_ci/dummy_image.py
+78 -72
View File
@@ -11,7 +11,7 @@ import zipfile
import tomlkit
sys.path.append(os.path.join(sys.path[0], ".."))
sys.path.append(os.path.join(sys.path[0], '..'))
from avbroot import external, ota, util
from avbroot.main import PARTITION_PRIORITIES, PATH_PAYLOAD
import ota_utils
@@ -27,12 +27,12 @@ def strtobool(val):
'val' is anything else.
"""
val = val.lower()
if val in ("y", "yes", "t", "true", "on", "1"):
if val in ('y', 'yes', 't', 'true', 'on', '1'):
return 1
elif val in ("n", "no", "f", "false", "off", "0"):
elif val in ('n', 'no', 'f', 'false', 'off', '0'):
return 0
else:
raise ValueError("invalid truth value %r" % (val,))
raise ValueError('invalid truth value %r' % (val, ))
class Crc32Hasher:
@@ -46,15 +46,15 @@ class Crc32Hasher:
def update(self, *args, **kwargs):
if self.existing_crc32:
args = args + (self.existing_crc32,)
args = args + (self.existing_crc32, )
self.existing_crc32 = crc32(*args, **kwargs)
@contextlib.contextmanager
def output_open(filename, mode, compress=True):
if compress and not filename.endswith(".gz"):
filename += ".gz"
if compress and not filename.endswith('.gz'):
filename += '.gz'
if compress:
with gzip.GzipFile(filename, mode, compresslevel=1, mtime=0) as f:
@@ -78,10 +78,10 @@ def output_writer(f_in, f_out, size, db=None):
crc_hash = Crc32Hasher()
# Combine sections if adjacent
if db and db[-1]["range"][1] == f_out.tell() - 1:
if db and db[-1]['range'][1] == f_out.tell() - 1:
entry = db.pop()
crc_hash.set_checksum(entry["checksum"])
section_start = entry["range"][0]
crc_hash.set_checksum(entry['checksum'])
section_start = entry['range'][0]
else:
section_start = f_out.tell()
@@ -90,45 +90,47 @@ def output_writer(f_in, f_out, size, db=None):
util.copyfileobj_n(f_in, f_out, size, hasher=crc_hash)
if mapping:
db.append(
{
"range": [section_start, section_end],
"checksum": crc_hash.get_checksum(),
}
)
db.append({
'range': [section_start, section_end],
'checksum': crc_hash.get_checksum(),
})
def download_file(f_out, url, section):
crc_hash = Crc32Hasher()
req = urllib.request.Request(url)
req.add_header("Range", f"bytes={section['range'][0]}-{section['range'][1]}")
req.add_header('Range',
f"bytes={section['range'][0]}-{section['range'][1]}")
with urllib.request.urlopen(req) as data:
util.copyfileobj_n(
data, f_out, section["range"][1] - section["range"][0] + 1, hasher=crc_hash
)
util.copyfileobj_n(data,
f_out,
section['range'][1] - section['range'][0] + 1,
hasher=crc_hash)
if crc_hash.get_checksum() != int(section["checksum"], 16):
raise Exception(f"ERROR: Checksum of range {section['range']} doesn't match")
if crc_hash.get_checksum() != int(section['checksum'], 16):
raise Exception(
f"ERROR: Checksum of range {section['range']} doesn't match")
def download_dummy_image(args):
with open(args.db, "r") as f:
with open(args.db, 'r') as f:
database = f.read()
database = tomlkit.parse(database)
device_entry = database["device"].get(args.device_name)
device_entry = database['device'].get(args.device_name)
if not device_entry:
raise Exception(f"ERROR: {args.device_name} cannot be found in the database")
raise Exception(
f'ERROR: {args.device_name} cannot be found in the database')
# Assume we always download the beginning and end (which makes sense for a ZIP file)
previous_offset = 0
with output_open(args.output, "wb", args.compress) as f_out:
for section in device_entry["sections"]:
util.zero_n(f_out, section["range"][0] - previous_offset)
download_file(f_out, device_entry["url"], section)
previous_offset = section["range"][1] + 1
with output_open(args.output, 'wb', args.compress) as f_out:
for section in device_entry['sections']:
util.zero_n(f_out, section['range'][0] - previous_offset)
download_file(f_out, device_entry['url'], section)
previous_offset = section['range'][1] + 1
def add_entry(db, db_entry, device_name, original_filename):
@@ -139,27 +141,27 @@ def add_entry(db, db_entry, device_name, original_filename):
entry = tomlkit.inline_table()
for key, value in section.items():
entry[key] = f"{value:x}" if key == "checksum" else value
entry[key] = f'{value:x}' if key == 'checksum' else value
new_sections.append(entry)
device_entry = {
device_name: {
"filename": os.path.basename(original_filename),
"url": "",
"sections": new_sections,
'filename': os.path.basename(original_filename),
'url': '',
'sections': new_sections,
}
}
with open(db, "r") as f:
with open(db, 'r') as f:
toml_db = tomlkit.load(f)
if toml_db.get("device").get(device_name) is not None:
if toml_db.get('device').get(device_name) is not None:
overwrite_entry = None
while overwrite_entry is None:
choice = input(
f"{device_name} is already in the database. Do you want to overwrite it? [y/n] "
f'{device_name} is already in the database. Do you want to overwrite it? [y/n] '
).lower()
try:
@@ -168,23 +170,25 @@ def add_entry(db, db_entry, device_name, original_filename):
pass
if not overwrite_entry:
print("Entry not added to the database. It is printed below instead:")
print(tomlkit.dumps({"device": device_entry}))
print(
'Entry not added to the database. It is printed below instead:'
)
print(tomlkit.dumps({'device': device_entry}))
return
toml_db.get("device").update(device_entry)
toml_db.get('device').update(device_entry)
with open(db, "w") as f:
with open(db, 'w') as f:
tomlkit.dump(toml_db, f)
print(f"{device_name} is added to the database")
print(f'{device_name} is added to the database')
def convert_image(args):
with zipfile.ZipFile(args.input, "r") as z:
with zipfile.ZipFile(args.input, 'r') as z:
info = z.getinfo(PATH_PAYLOAD)
with z.open(info, "r") as f:
with z.open(info, 'r') as f:
_, manifest, blob_offset = ota.parse_payload(f)
file_offset, _ = ota_utils.GetZipEntryOffset(z, info)
@@ -203,28 +207,27 @@ def convert_image(args):
continue
if old_data_offset > op.data_offset:
raise Exception("Operations are expected to be ordered")
raise Exception('Operations are expected to be ordered')
old_data_offset = op.data_offset
partition_operations[p.partition_name] = p.operations
db_entry = []
with open(args.input, "rb") as f_in, output_open(
args.output, "wb", args.compress
) as f_out:
with open(args.input, 'rb') as f_in, output_open(args.output, 'wb',
args.compress) as f_out:
# Copy zip data until payload.bin + payload.bin header
output_writer(f_in, f_out, file_offset + blob_offset, db_entry)
# Sort partitions based on data_offset of first operation and iterate
for _, p_operations in sorted(
partition_operations.items(), key=lambda x: x[1][0].data_offset
):
for _, p_operations in sorted(partition_operations.items(),
key=lambda x: x[1][0].data_offset):
# Copy all data until first partition offset
output_writer(
f_in,
f_out,
p_operations[0].data_offset - (f_in.tell() - file_offset - blob_offset),
p_operations[0].data_offset -
(f_in.tell() - file_offset - blob_offset),
db_entry,
)
@@ -232,8 +235,9 @@ def convert_image(args):
if op.type == Type.ZERO or op.type == Type.DISCARD:
continue
if f_out.tell() != (file_offset + blob_offset + op.data_offset):
raise Exception("f_out should be equal to the offset")
if f_out.tell() != (file_offset + blob_offset +
op.data_offset):
raise Exception('f_out should be equal to the offset')
output_writer(None, f_out, op.data_length, db_entry)
f_in.seek(f_out.tell())
@@ -248,29 +252,31 @@ def convert_image(args):
def parse_args(args=None):
parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers(
dest="subcommand", required=True, help="Subcommands"
)
subparsers = parser.add_subparsers(dest='subcommand',
required=True,
help='Subcommands')
base = argparse.ArgumentParser(add_help=False)
base.add_argument(
"--compress",
'--compress',
action=argparse.BooleanOptionalAction,
default=True,
help="Compress output file",
help='Compress output file',
)
base.add_argument("--db", required=True, help="Path to database file")
base.add_argument("--output", required=True, help="Path to new OTA zip")
base.add_argument("device_name", help="Name of device (as in database)")
base.add_argument('--db', required=True, help='Path to database file')
base.add_argument('--output', required=True, help='Path to new OTA zip')
base.add_argument('device_name', help='Name of device (as in database)')
convert = subparsers.add_parser(
"convert", help="Convert full OTA zip to dummy zip", parents=[base]
)
convert.add_argument("--input", required=True, help="Path to original OTA zip")
convert = subparsers.add_parser('convert',
help='Convert full OTA zip to dummy zip',
parents=[base])
convert.add_argument('--input',
required=True,
help='Path to original OTA zip')
subparsers.add_parser(
"download", help="Assemble dummy zip while downloading", parents=[base]
)
subparsers.add_parser('download',
help='Assemble dummy zip while downloading',
parents=[base])
return parser.parse_args(args=args)
@@ -278,13 +284,13 @@ def parse_args(args=None):
def main(args=None):
args = parse_args(args=args)
if args.subcommand == "convert":
if args.subcommand == 'convert':
convert_image(args)
elif args.subcommand == "download":
elif args.subcommand == 'download':
download_dummy_image(args)
else:
raise NotImplementedError()
if __name__ == "__main__":
if __name__ == '__main__':
main()
tests_ci/test_device.py
+59 -63
View File
@@ -10,23 +10,24 @@ import urllib
import tomlkit
sys.path.append(os.path.join(sys.path[0], ".."))
sys.path.append(os.path.join(sys.path[0], '..'))
import avbroot.main
from avbroot import util
import dummy_image
def test_device(test_file, magisk_file, hashes, workdir, no_test, db, device_name):
output_file = os.path.join(workdir, "output.zip")
test_key_prefix = os.path.join(
sys.path[0], os.pardir, "tests", "keys", "TEST_KEY_DO_NOT_USE_"
)
def test_device(test_file, magisk_file, hashes, workdir, no_test, db,
device_name):
output_file = os.path.join(workdir, 'output.zip')
test_key_prefix = os.path.join(sys.path[0], os.pardir, 'tests', 'keys',
'TEST_KEY_DO_NOT_USE_')
# We intentionally create a zip file with an invalid CRC for the payload.
# The CRC of the zip will be checked when EOF is reached.
# Although we never intentionally read all the way to the end, prefetching of a file
# might lead to read until EOF and triggering the CRC check, which will fail.
with unittest.mock.patch('zipfile.ZipExtFile._update_crc', lambda _, __: None):
with unittest.mock.patch('zipfile.ZipExtFile._update_crc',
lambda _, __: None):
avbroot.main.main([
'patch',
'--input',
@@ -43,42 +44,40 @@ def test_device(test_file, magisk_file, hashes, workdir, no_test, db, device_nam
test_key_prefix + 'ota.crt',
])
avbroot.main.main(
[
"extract",
"--input",
output_file,
"--directory",
workdir,
]
)
avbroot.main.main([
'extract',
'--input',
output_file,
'--directory',
workdir,
])
if hashes is None:
with open(db, "r") as f:
with open(db, 'r') as f:
toml_db = tomlkit.load(f)
new_hashes = tomlkit.table()
with open(os.path.join(workdir, "output.zip"), "rb") as f:
with open(os.path.join(workdir, 'output.zip'), 'rb') as f:
new_hashes.update(
{"output.zip": util.hash_file(f, hashlib.md5()).hexdigest()}
)
{'output.zip': util.hash_file(f, hashlib.md5()).hexdigest()})
for i in sorted(glob.glob("*.img", root_dir=workdir)):
with open(os.path.join(workdir, i), "rb") as f:
new_hashes.update({i: util.hash_file(f, hashlib.md5()).hexdigest()})
for i in sorted(glob.glob('*.img', root_dir=workdir)):
with open(os.path.join(workdir, i), 'rb') as f:
new_hashes.update(
{i: util.hash_file(f, hashlib.md5()).hexdigest()})
toml_db.get("device").get(device_name).append("hashes", new_hashes)
toml_db.get('device').get(device_name).append('hashes', new_hashes)
with open(db, "w") as f:
with open(db, 'w') as f:
tomlkit.dump(toml_db, f)
print(f"{device_name} hashes are added to the database")
print(f'{device_name} hashes are added to the database')
hashes = dict(new_hashes)
if not no_test:
for filename, md5_hash in hashes.items():
with open(os.path.join(workdir, filename), "rb") as f:
with open(os.path.join(workdir, filename), 'rb') as f:
assert util.hash_file(f, hashlib.md5()).hexdigest() == md5_hash
os.remove(os.path.join(workdir, filename))
@@ -88,7 +87,7 @@ def download_magisk(f_out, url, checksum):
req = urllib.request.Request(url)
with urllib.request.urlopen(req) as d:
size = d.info()["Content-Length"]
size = d.info()['Content-Length']
util.copyfileobj_n(d, f_out, int(size), hasher=sha256_hash)
if sha256_hash.hexdigest() != checksum:
@@ -98,19 +97,19 @@ def download_magisk(f_out, url, checksum):
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"-d",
"--device",
action="append",
'-d',
'--device',
action='append',
required=True,
help="Device name to test against",
)
parser.add_argument("--db", required=True, help="Path to database file")
parser.add_argument(
"--no-test", action="store_true", help="Don't test and clean in final steps"
)
parser.add_argument(
"--workdir", required=True, help="Path to directory which contains images"
help='Device name to test against',
)
parser.add_argument('--db', required=True, help='Path to database file')
parser.add_argument('--no-test',
action='store_true',
help="Don't test and clean in final steps")
parser.add_argument('--workdir',
required=True,
help='Path to directory which contains images')
return parser.parse_args()
@@ -118,45 +117,42 @@ def parse_args():
def test_main():
args = parse_args()
with open(args.db, "r") as f:
with open(args.db, 'r') as f:
database = tomlkit.load(f)
magisk_file = os.path.join(args.workdir, database["magisk"]["filename"])
magisk_file = os.path.join(args.workdir, database['magisk']['filename'])
if not os.path.isfile(magisk_file):
with open(magisk_file, "wb") as f_out:
print("Downloading Magisk...")
download_magisk(
f_out, database["magisk"]["url"], database["magisk"]["sha256"]
)
with open(magisk_file, 'wb') as f_out:
print('Downloading Magisk...')
download_magisk(f_out, database['magisk']['url'],
database['magisk']['sha256'])
for device_name in args.device:
device_entry = database.get("device").get(device_name)
device_entry = database.get('device').get(device_name)
if not device_entry:
raise Exception(f"Device {device_name} not found in database")
raise Exception(f'Device {device_name} not found in database')
test_file = os.path.join(args.workdir, device_entry["filename"])
test_file = os.path.join(args.workdir, device_entry['filename'])
if not os.path.isfile(test_file):
print(f"No input file found for {device_name}. Downloading now...")
print(f'No input file found for {device_name}. Downloading now...')
dummy_image.main(
[
"download",
"--no-compress",
"--db",
args.db,
"--output",
test_file,
device_name,
]
)
dummy_image.main([
'download',
'--no-compress',
'--db',
args.db,
'--output',
test_file,
device_name,
])
test_device(
test_file,
magisk_file,
device_entry.get("hashes"),
device_entry.get('hashes'),
args.workdir,
args.no_test,
args.db,
@@ -164,5 +160,5 @@ def test_main():
)
if __name__ == "__main__":
if __name__ == '__main__':
test_main()