librem5-flash-image 15.2 KB
Newer Older
1
#!/usr/bin/env python3
2
3
4
#
# SPDX-License-Identifier: GPL-3.0-or-later
# Copyright (C) 2018-2021 Purism SPC
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

import argparse
import datetime
import hashlib
import itertools
import logging
import lzma
import os
import re
import requests
import shutil
import subprocess
import sys
import tempfile
import time
import tqdm
import yaml

23
24
25
26
27
28
try:
    import jenkins
except ImportError:
    logging.error("Cannot load Jenkins\nTry running: sudo apt install python3-jenkins")
    exit(1)

29
30
31
32
33
34
35
36
37
38
try:
    import coloredlogs
    have_colored_logs = True
except ImportError:
    have_colored_logs = False


from urllib.parse import urljoin

JENKINS = 'https://arm01.puri.sm'
39
BOARD_TYPE = 'librem5r4'
40
41
VALID_PHONE_BOARD_TYPES = ['librem5r2', 'librem5r3', 'librem5r4']
VALID_DEVKIT_BOARD_TYPES = ['devkit']
Guido Gunther's avatar
Guido Gunther committed
42
DIST = 'byzantium'
43
IMAGE = '{}.img'
44
IMAGE_VARIANT = 'plain'
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
META_YAML = 'files/meta.yml'
IMAGE_JOB_NAME = 'Images/Image Build'
UBOOT = 'u-boot-{}.imx'
UBOOT_JOB_NAME = 'u-boot_builds/uboot_{}_build'
UUU_SCRIPT = 'flash_{}.lst'
UUU_SCRIPT_TMPL = '''uuu_version 1.0.1

{mods}
SDP: boot -f {uboot}
# This command will be run when use SPL
SDPU: delay 1000
SDPU: write -f {uboot} -offset 0x57c00
SDPU: jump
SDPV: delay 1000
SDPV: write -f {uboot} -skipspl
SDPV: 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
Guido Gunther's avatar
Guido Gunther committed
71
RETRIES = 10
72

73
74
75
76
77
78
79
80
81
82
UDEV_RULES = '''
SUBSYSTEM!="usb", GOTO="librem5_devkit_rules_end"
# Devkit USB flash
ATTR{idVendor}=="1fc9", ATTR{idProduct}=="012b", GROUP+="plugdev", TAG+="uaccess"
ATTR{idVendor}=="0525", ATTR{idProduct}=="a4a5", GROUP+="plugdev", TAG+="uaccess"
ATTR{idVendor}=="0525", ATTR{idProduct}=="b4a4", GROUP+="plugdev", TAG+="uaccess"
ATTR{idVendor}=="316d", ATTR{idProduct}=="4c05", GROUP+="plugdev", TAG+="uaccess"
LABEL="librem5_devkit_rules_end"
'''

83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271

class VerifyImageException(Exception):
    pass


class PrematureEndException(Exception):
    pass


def verify_image(image, meta):
    m = hashlib.sha256()
    size = int(meta['image']['size'])
    hexdigest = meta['image']['sha256sum']

    filesize = os.path.getsize(image)
    if filesize != size:
        raise VerifyImageException(
            "Image file \"{}\" size {} does not match {}".format(
                os.path.basename(image), filesize, size))

    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 resuming_stream(url, expected_size, max_attempts):
    position = 0

    if max_attempts < 1:
        retries = itertools.count()
    else:
        retries = range(max_attempts)

    for i in retries:
        try:
            resp = requests.get(url,
                                stream=True,
                                headers={'Range': 'bytes={}-'.format(position)}
                                )
            resp.raise_for_status()

            if resp.status_code != requests.codes.partial_content:
                position = 0
            logging.debug('Proceeding from {} bytes'.format(position))

            for data in resp.iter_content(BLOCK_SIZE):
                position += len(data)
                yield data

            if position < expected_size:
                raise PrematureEndException()
            return
        except (requests.exceptions.ConnectionError, PrematureEndException):
            if i == max_attempts - 1:
                logging.error("Max connection errors reached, aborting")
                raise
            logging.info("Connection error, retrying")
            time.sleep(5)


def stream_file(url, attempts):
    resp = requests.head(url, stream=True)
    resp.raise_for_status()
    ts = int(resp.headers.get('content-length', 0))
    return resuming_stream(url, ts, attempts), ts


def needs_download(target, meta):
    if not os.path.exists(target):
        return True

    try:
        verify_image(target, meta)
    except VerifyImageException:
        return True

    return False


def download_image(url, target, attempts):
    decomp = lzma.LZMADecompressor()

    # 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)
    resp.raise_for_status()
    meta = yaml.safe_load(resp.text)
    uncompressed_size = int(meta['image']['size'])
    logging.debug("Image size is %d", uncompressed_size)

    if not needs_download(target, meta):
        logging.info("Image already up to date - no download needed.")
        return

    logging.info("Downloading image from {}".format(url))
    stream, ts = stream_file(url, attempts)
    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 stream:
            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()
    verify_image(target, meta)


def find_image(jobname, type, variant, dist):
    server = jenkins.Jenkins(JENKINS)
    logging.info("Looking for {} {} {} image".format(type, variant, dist))
    try:
        info = server.get_job_info(jobname)
    except jenkins.NotFoundException:
        logging.error("Job %s not found", jobname)
        return None
    for build in info['builds']:
        resp = requests.get(build['url'] + '/api/json')
        resp.raise_for_status()
        json = resp.json()
        if (json['description'].startswith(variant + ' ' + 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)

    try:
        info = server.get_job_info(jobname)
    except jenkins.NotFoundException:
        logging.error("Job %s not found", jobname)
        return None
    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, mods):
    with open(target, 'w+') as f:
        f.write(UUU_SCRIPT_TMPL.format(image=os.path.basename(image),
                                       uboot=os.path.basename(uboot),
                                       mods=mods))


272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
def wait_for_librem(librem_msg="default"):
    found = False
    printed_msg = False

    if librem_msg == "default":
        librem_msg = """
            Turn all Hardware-Kill-Switches off
            Remove battery
            Hold volume-up
            Insert the USB-c cable: (red light blinks, no green light)
            Reinsert the battery: (red light is constantly on, the script will continue)\n\nSearching..."""

    while not found:
        returned_text = subprocess.check_output("lsusb", shell=True, universal_newlines=True)
        for line in returned_text.split("\n"):
            if "NXP Semiconductors" in line:
                found = True
        if not found and not printed_msg:
            print(librem_msg)
            printed_msg = True
        time.sleep(1)


295
296
297
298
299
300
301
def flash_image(uuu_target, debug):
    if debug:
        subprocess.check_call(['uuu', '-v', uuu_target])
    else:
        subprocess.check_call(['uuu', uuu_target])


302
303
304
305
306
307
def get_board(path):
    """
    Get board type from script name.

    In case of doubt return the default board.

308
309
    >>> get_board("/a/b/librem5-flash-image")
    'librem5r4'
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
    >>> get_board("/a/b/librem5-devkit-flash-image")
    'devkit'
    >>> get_board("/a/b/librem5r2-flash-image")
    'librem5r2'
    >>> get_board("/a/b/librem5r3-flash-image")
    'librem5r3'
    >>> get_board("/a/b/librem5r4-flash-image")
    'librem5r4'
    >>> get_board("/a/b/schwobel")
    'librem5r4'
    """
    name = os.path.split(path)[-1]
    parts = name.split('-')
    if len(parts) < 3 or len(parts) > 4:
        return BOARD_TYPE

    if parts[-2:] != ['flash', 'image']:
        return BOARD_TYPE

    parts = parts[:-2]
330
331
332
333
334
    if len(parts) == 1:
        if parts[0] in VALID_PHONE_BOARD_TYPES:
            return parts[0]
        else:
            return BOARD_TYPE
335
336
337
338
339
340
341
342

    if parts[0] == 'librem5' and parts[1] in VALID_DEVKIT_BOARD_TYPES:
        return parts[1]

    # People do odd things so use default in this case
    return BOARD_TYPE


343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
def setup_udev():
    rules = "/etc/udev/rules.d/70-librem5-flash-image.rules"
    if os.path.exists(rules):
        logging.info(f"Udev rules '{rules}' already present")
        return 0

    try:
        with open(rules, 'w+') as f:
            f.write(UDEV_RULES)
    except PermissionError as e:
        logging.error(f"{e} - are you root?")
        return 1

    subprocess.check_call(['udevadm', 'control', '--reload'])
    logging.info(f"Udev rules '{rules}' updated. You can now flash without root permissions.")
    return 0


361
362
def main():
    uuu_mods = ''
Guido Gunther's avatar
Guido Gunther committed
363
    parser = argparse.ArgumentParser(description='Flash a Librem 5 or Librem 5 Devkit.')
364
365
366
367
368
369
370
371
    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')
    parser.add_argument('--skip-flash', action='store_true', default=False,
                        help='Do all the preparations but don\'t flash')
Guido Gunther's avatar
Guido Gunther committed
372
    parser.add_argument('--download-attempts', type=int, default=RETRIES,
373
                        help="Maximum number of attempts to resume "
Guido Gunther's avatar
Guido Gunther committed
374
                        "image download. 0: unlimited, default is {}".format(RETRIES))
375
376
377
    parser.add_argument('--variant', choices=['plain', 'luks'], default=IMAGE_VARIANT,
                        help="Image variant to download ( plain, luks ), "
                        "default is '{}'".format(IMAGE_VARIANT))
378
    parser.add_argument('--board', choices=VALID_DEVKIT_BOARD_TYPES + VALID_PHONE_BOARD_TYPES, default=None,
379
                        help="Type of the board to download ( devkit, librem5r2, librem5r3, librem5r4 ) "
380
                        "default is 'auto'")
381
382
383
384
385
386
387
388

    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,
                       help='Jenkins job to download the uboot from')
389
390
391
392

    group = parser.add_argument_group(title='Initial setup options')
    group.add_argument('--udev', action="store_true", default=False,
                       help='Setup udev rules to flash without root permissions')
393
394
395
396
397
398
399
400
    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')

401
402
403
    if args.udev:
        return setup_udev()

404
405
    board = args.board if args.board else get_board(sys.argv[0])

406
407
    # Check available downloads upfront so it's less likely we fail
    # later:
408
    image_ref = find_image(args.image_job, board, args.variant, args.dist)
409
410
411
412
413
414
415
416
417
    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

    if args.uboot_job:
        uboot_ref = find_uboot(args.uboot_job)
418
        uboot_board = board
419
420
    else:
        # uboot builds don't carry board revisions (yet?)
421
        uboot_board = board[:-2] if re.match('librem5r[0-9]$', board) else board
422
423
424
425
426
427
428
429
430
        uboot_ref = find_uboot(UBOOT_JOB_NAME.format(uboot_board))

    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

431
    if board in ["librem5r3", 'librem5r4']:
432
433
434
435
436
437
438
439
        uuu_mods = "CFG: SDP: -chip MX8MQ -compatible MX8MQ -vid 0x316d -pid 0x4c05"

    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)

440
        image_target = os.path.join(outdir, IMAGE.format(board))
441
        uboot_target = os.path.join(outdir, UBOOT.format(uboot_board))
442
        uuu_target = os.path.join(outdir, UUU_SCRIPT.format(board))
443

444
        download_image(urljoin(image_ref['url'], 'artifact/{}.xz').format(IMAGE.format(board)),
445
446
447
448
449
450
                       image_target, args.download_attempts)
        download_uboot(urljoin(uboot_ref['url'],
                       'artifact/output/uboot-{}/{}'.format(uboot_board, UBOOT.format(uboot_board))),
                       uboot_target)
        write_uuu_script(uuu_target, image_target, uboot_target, uuu_mods)
        if not args.skip_flash:
451
452
453
454
            if board == "devkit":
                wait_for_librem(librem_msg="Please connect your devkit now.\nSearching...")
            else:
                wait_for_librem()
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
            flash_image(uuu_target, args.debug)
    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())