1# Copyright (c) 2023 Peter Johanson <peter@peterjohanson.com>
2#
3# SPDX-License-Identifier: Apache-2.0
4
5'''UF2 runner (flash only) for UF2 compatible bootloaders.'''
6
7from pathlib import Path
8from shutil import copy
9
10from runners.core import ZephyrBinaryRunner, RunnerCaps
11
12try:
13    import psutil  # pylint: disable=unused-import
14    MISSING_PSUTIL = False
15except ImportError:
16    # This can happen when building the documentation for the
17    # runners package if psutil is not on sys.path. This is fine
18    # to ignore in that case.
19    MISSING_PSUTIL = True
20
21class UF2BinaryRunner(ZephyrBinaryRunner):
22    '''Runner front-end for copying to UF2 USB-MSC mounts.'''
23
24    def __init__(self, cfg, board_id=None):
25        super().__init__(cfg)
26        self.board_id = board_id
27
28    @classmethod
29    def name(cls):
30        return 'uf2'
31
32    @classmethod
33    def capabilities(cls):
34        return RunnerCaps(commands={'flash'})
35
36    @classmethod
37    def do_add_parser(cls, parser):
38        parser.add_argument('--board-id', dest='board_id',
39                            help='Board-ID value to match from INFO_UF2.TXT')
40
41    @classmethod
42    def do_create(cls, cfg, args):
43        return UF2BinaryRunner(cfg, board_id=args.board_id)
44
45    @staticmethod
46    def get_uf2_info_path(part) -> Path:
47        return Path(part.mountpoint) / "INFO_UF2.TXT"
48
49    @staticmethod
50    def is_uf2_partition(part):
51        try:
52            return ((part.fstype in ['vfat', 'FAT', 'msdos']) and
53                    UF2BinaryRunner.get_uf2_info_path(part).is_file())
54        except PermissionError:
55            return False
56
57    @staticmethod
58    def get_uf2_info(part):
59        lines = UF2BinaryRunner.get_uf2_info_path(part).read_text().splitlines()
60
61        lines = lines[1:] # Skip the first summary line
62
63        def split_uf2_info(line: str):
64            k, _, val = line.partition(':')
65            return k.strip(), val.strip()
66
67        return {k: v for k, v in (split_uf2_info(line) for line in lines) if k and v}
68
69    def match_board_id(self, part):
70        info = self.get_uf2_info(part)
71
72        return info.get('Board-ID') == self.board_id
73
74    def get_uf2_partitions(self):
75        parts = [part for part in psutil.disk_partitions() if self.is_uf2_partition(part)]
76
77        if (self.board_id is not None) and parts:
78            parts = [part for part in parts if self.match_board_id(part)]
79            if not parts:
80                self.logger.warning("Discovered UF2 partitions don't match Board-ID '%s'",
81                                    self.board_id)
82
83        return parts
84
85    def copy_uf2_to_partition(self, part):
86        self.ensure_output('uf2')
87
88        copy(self.cfg.uf2_file, part.mountpoint)
89
90    def do_run(self, command, **kwargs):
91        if MISSING_PSUTIL:
92            raise RuntimeError(
93                'could not import psutil; something may be wrong with the '
94                'python environment')
95
96        partitions = self.get_uf2_partitions()
97        if not partitions:
98            raise RuntimeError('No matching UF2 partitions found')
99
100        if len(partitions) > 1:
101            raise RuntimeError('More than one matching UF2 partitions found')
102
103        part = partitions[0]
104        self.logger.info("Copying UF2 file to '%s'", part.mountpoint)
105        self.copy_uf2_to_partition(part)
106