Apply patches to additional zephyr modules

apply_patches.py

#!/usr/bin/env python3

# Copyright (c) 2024, Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0

"""Apply patches to one or more Zephyr modules

Usage:

west patch [-b DIR] [-l FILE] [-w DIR] [-m NAME] [-r]

options:
  -b DIR, --patch-base DIR
                        Directory containing patch files
  -l FILE, --patch-yml FILE
                        Path to patches.yml file
  -w DIR, --west-workspace DIR
                        West workspace
  -m NAME, --module NAME
                        Only apply patches for module
  -r, --roll-back       Roll back if any patch fails to apply

```
patches:
  - path: zephyr/too-many-songs-about-rainbows.patch
    sha256sum: ccd2741eefec156387944177546393d76f9dcb8395af8222dcb2d71ab740808b
    module: zephyr
    author: Kermit D. Frog
    email: [email protected]
    date: 2024-05-15
    upstreamable: True
    merge-pr: https://github.com/zephyrproject-rtos/zephyr/pull/123456
    issue: {}
    merge-status: False
    merge-commit: {}
    merge-date: 2024-05-16
    apply-command: {}
    comments: |
      Note: Why are there so many??
```
"""

import argparse
import hashlib
import os
import subprocess
import sys
import yaml

from pathlib import Path

# TODO: eventually convert this to a WestCommand
# from west.commands import WestCommand


class Patcher(object):
    def __init__(self, patch_base, west_workspace, yml, roll_back):
        self._patch_base = patch_base
        self._west_workspace = west_workspace
        self._yml = yml
        self._roll_back = roll_back

    def list(self):
        for patch_info in self._yml['patches']:
            print(patch_info)

    def apply(self):
        failed_patch = None
        cleanup_mods = {}

        if not (self._yml and 'patches' in self._yml):
            print(f'no patches to apply')
            return

        for patch_info in self._yml['patches']:

            pth = patch_info['path']
            patch_path = os.path.realpath(Path(self._patch_base) / pth)

            print(f'reading patch file {pth}')
            patch_file_data = None

            try:
                with open(patch_path, 'rb') as pf:
                    patch_file_data = pf.read()
            except Exception as e:
                print(f'failed to read {pth}: {e}')
                failed_patch = pth
                break

            print('checking integrity... ', end='')
            expect_sha256 = patch_info['sha256sum']
            hasher = hashlib.sha256()
            hasher.update(patch_file_data)
            actual_sha256 = hasher.hexdigest()
            if actual_sha256 != expect_sha256:
                print('BAD')
                print(f'sha256 mismatch for {pth}:\nexpect: {expect_sha256}\nactual: {actual_sha256}')
                failed_patch = pth
                break
            print('GOOD')
            patch_file_data = None

            mod = patch_info['module']
            mod_path = os.path.realpath(Path(self._west_workspace) / mod)
            cleanup_mods[mod] = mod_path
            print(f'attempting to patch {mod}')

            origdir = os.getcwd()
            os.chdir(mod_path)
            proc = subprocess.run(['git', 'apply', patch_path])
            if proc.returncode:
                print(f'failed to apply patch: returncode {proc.returncode}')
                print(proc.stderr)
                failed_patch = pth
            os.chdir(origdir)

        if not failed_patch:
            print('all patches applied successfully \\o/')
            return

        if self._roll_back:
            for mod, mod_path in cleanup_mods.items():
                print(f'rolling back changes to {mod}')
                os.chdir(mod_path)
                subprocess.run(['git', 'checkout', '.'])
                subprocess.run(['git', 'clean', '-x', '-d', '-f'])

        raise ValueError(f'failed to apply patch {pth}')


def parse_args():
    parser = argparse.ArgumentParser(allow_abbrev=False)
    parser.add_argument('-b', '--patch-base', help='Directory containing patch files',
                        metavar='DIR', default=Path('zephyr/patches'))
    parser.add_argument('-l', '--patch-yml', help='Path to patches.yml file',
                        metavar='FILE', default=Path('zephyr/patches.yml'))
    parser.add_argument('-w', '--west-workspace',
                        help='West workspace', metavar='DIR')
    parser.add_argument('-r', '--roll-back', help='Roll back if any patch fails to apply',
                        action='store_true', default=False)

    args = parser.parse_args()

    if not os.path.isdir(args.patch_base):
        raise FileNotFoundError(args.patch_base)
    if not os.path.isfile(args.patch_yml):
        raise FileNotFoundError(args.patch_yml)

    west_config = Path(f'{args.west_workspace}/.west/config')
    if not os.path.isfile(west_config):
        raise FileNotFoundError(west_config)

    with open(args.patch_yml) as f:
        yml = yaml.safe_load(f)

    setattr(args, 'yml', yml)

    return args


def main():
    args = parse_args()

    if not args:
        return os.EX_DATAERR

    p = Patcher(args.patch_base, args.west_workspace, args.yml, args.roll_back)
    p.apply()

    return os.EX_OK


if __name__ == '__main__':
    sys.exit(main())