#!/usr/bin/python3 import argparse import datetime import hashlib import jenkins import logging import lzma import os import requests import shutil import subprocess import sys import tempfile import tqdm import yaml try: import coloredlogs have_colored_logs = True except ImportError: have_colored_logs = False from urllib.parse import urljoin JENKINS = 'https://arm01.puri.sm' TYPE = 'devkit' DIST = 'buster+ci' IMAGE = 'devkit.img' META_YAML = 'files/meta.yml' IMAGE_JOB_NAME = 'Images/Image Build' UBOOT = 'u-boot-devkit-recovery.imx' UBOOT_JOB_NAME = 'u-boot_builds/devkit-recovery_uboot_build' UUU_SCRIPT = 'flash_devkit.lst' UUU_SCRIPT_TMPL = '''uuu_version 1.0.1 SDP: boot -f {uboot} # This command will be run when use SPL SDPU: delay 1000 SDPU: write -f {uboot} -offset 0x57c00 SDPU: jump # This command will be run when ROM support stream mode SDPS: boot -f {uboot} SDPU: delay 1000 #FB: ucmd setenv fastboot_buffer 0x43000000 FB: ucmd setenv fastboot_dev mmc FB: ucmd setenv mmcdev 0 FB: flash -raw2sparse all {image} FB: Done ''' BLOCK_SIZE = 8192 UNCOMPRESSED_SIZE = 3600000000 class VerifyImageException(Exception): pass def verify_image(image, meta): m = hashlib.sha256() size = int(meta['image']['size']) hexdigest = meta['image']['sha256sum'] logging.info("Calculating sha256sum of {}".format(image)) bar = tqdm.tqdm(total=size, desc='Checking', leave=False) with open(image, 'rb') as f: while True: data = f.read(BLOCK_SIZE) if data: m.update(data) bar.update(n=len(data)) else: break bar.close() if m.hexdigest() != hexdigest: raise VerifyImageException("Checksum of image {} " "does not match {}".format(m.hexdigest(), hexdigest)) def download_image(url, target): decomp = lzma.LZMADecompressor() logging.info("Downloading image from {}".format(url)) # We expect metadata to be right next to the image meta_yml_url = "{}/{}".format(url.rsplit('/', 1)[0], META_YAML) resp = requests.get(meta_yml_url) # No meta data is no yet fatal try: resp.raise_for_status() meta = yaml.safe_load(resp.text) uncompressed_size = int(meta['image']['size']) logging.debug("Image size is %d", uncompressed_size) except requests.exceptions.HTTPError: logging.warning("No meta data found, skipping integrity check") meta = None uncompressed_size = UNCOMPRESSED_SIZE resp = requests.get(url, stream=True) resp.raise_for_status() ts = int(resp.headers.get('content-length', 0)) download_bar = tqdm.tqdm(total=ts, desc='Download', leave=False) decompress_bar = tqdm.tqdm(total=uncompressed_size, desc='Decompr.', leave=False) with open(target, 'wb+') as f: for data in resp.iter_content(BLOCK_SIZE): if data: out = decomp.decompress(data) decompress_bar.update(len(out)) f.write(out) download_bar.update(n=len(data)) download_bar.close() decompress_bar.close() if meta: verify_image(target, meta) def find_image(jobname, type, dist): server = jenkins.Jenkins(JENKINS) logging.info("Looking for {} {} image".format(type, dist)) info = server.get_job_info(jobname) for build in info['builds']: resp = requests.get(build['url'] + '/api/json') resp.raise_for_status() json = resp.json() if (json['description'].startswith(type + ' ') and dist in json['description'] and json['result'] == 'SUCCESS'): found = json break else: found = None return found def find_uboot(jobname): server = jenkins.Jenkins(JENKINS) info = server.get_job_info(jobname) for build in info['builds']: resp = requests.get(build['url'] + '/api/json') resp.raise_for_status() json = resp.json() if (json['result'] == 'SUCCESS'): found = json break else: found = None return found def download_uboot(url, target): logging.info("Downloading uboot from {}".format(url)) resp = requests.get(url, stream=True) resp.raise_for_status() ts = int(resp.headers.get('content-length', 0)) bar = tqdm.tqdm(total=ts, leave=False) with open(target, 'wb+') as f: for data in resp.iter_content(BLOCK_SIZE): if data: f.write(data) bar.update(n=len(data)) def write_uuu_script(target, image, uboot): with open(target, 'w+') as f: f.write(UUU_SCRIPT_TMPL.format(image=os.path.basename(image), uboot=os.path.basename(uboot))) def flash_image(uuu_target): subprocess.check_call(['uuu', uuu_target]) def main(): parser = argparse.ArgumentParser(description='Process some integers.') parser.add_argument('--dir', type=str, default=None, help='Download files to dir (instead of a temporary directory)') parser.add_argument('--dist', type=str, default=DIST, help="Download an image for this distribution, default is '{}'".format(DIST)) parser.add_argument('--skip-cleanup', action='store_true', default=False, help='Skip temporary directory cleanup') group = parser.add_argument_group(title='Testing and debugging options') group.add_argument('--debug', action="store_true", default=False, help='Enable debug output') group.add_argument('--image-job', type=str, default=IMAGE_JOB_NAME, help='Jenkins job to download the image from') group.add_argument('--uboot-job', type=str, default=UBOOT_JOB_NAME, help='Jenkins job to download the uboot from') args = parser.parse_args() level = logging.DEBUG if args.debug else logging.INFO if have_colored_logs: coloredlogs.install(level=level, fmt='%(asctime)s %(levelname)s %(message)s') else: logging.basicConfig(level=level, format='%(asctime)s %(levelname)s %(message)s') # Check available downloads upfront so it's less likely we fail # later: image_ref = find_image(args.image_job, TYPE, args.dist) if image_ref: image_ref['ts'] = datetime.datetime.fromtimestamp(image_ref['timestamp'] / 1000).strftime('%c') logging.info("Found disk image Build {id} '{description}' from {ts}".format(**image_ref)) else: logging.error("No matching image found") return 1 uboot_ref = find_uboot(args.uboot_job) if uboot_ref: uboot_ref['ts'] = datetime.datetime.fromtimestamp(uboot_ref['timestamp'] / 1000).strftime('%c') logging.info("Found uboot Build {id} from {ts}".format(**uboot_ref)) else: logging.error("No matching uboot found") return 1 outdir = args.dir if args.dir is not None else tempfile.mkdtemp(prefix='devkit_image_', dir='.') try: logging.info("Downloading to {}".format(outdir)) if args.dir == outdir: os.makedirs(args.dir, exist_ok=True) image_target = os.path.join(outdir, IMAGE) uboot_target = os.path.join(outdir, UBOOT) uuu_target = os.path.join(outdir, UUU_SCRIPT) download_image(urljoin(image_ref['url'], 'artifact/{}.xz').format(IMAGE), image_target) download_uboot(urljoin(uboot_ref['url'], 'artifact/build/{}'.format(UBOOT)), uboot_target) write_uuu_script(uuu_target, image_target, uboot_target) flash_image(uuu_target) except VerifyImageException as e: logging.error(e) return 1 except KeyboardInterrupt: logging.error("CTRL-C pressed.") return 1 finally: if args.dir != outdir and not args.skip_cleanup: logging.info("Cleaning up.") shutil.rmtree(outdir) return 0 if __name__ == '__main__': sys.exit(main())