#!/usr/bin/env python3

# SPDX-FileCopyrightText: Copyright (c) 2020-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: MIT
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

import argparse
import os
import re
import shutil
import sys
import tempfile

EXTLINUX="/boot/extlinux/extlinux.conf"

# For addition, ignore the given param if it already exists in the current cmdline.
# Otherwise append it in the end as the last takes effect.
def update_kernel_params(current, additions, removals):
    # strip extra spaces if any
    current = re.sub(' +', ' ', current.rstrip())

    current_params = current.split()
    for param in additions:
        if param not in current_params:
            current_params.append(param)
    for param in removals:
        if param in current_params:
            current_params.remove(param)
    return ' '.join(current_params)

def add_extlinux_entry(image_path, label, mlabel, add_kernel_params, rm_kernel_params):
    # Check extlinux.conf
    if not os.path.exists(EXTLINUX):
        print("%s file not found" % EXTLINUX)
        sys.exit(1)

    if not os.access(EXTLINUX, os.R_OK | os.W_OK):
        print("%s no permission" % EXTLINUX)
        sys.exit(1)

    with open(EXTLINUX, 'r') as fin:
        contents = fin.readlines()

    new_entry = True
    default_label = None
    default_append = None
    current_label = None
    current_mlabel = None
    current_linux = None
    current_append = None
    out = []

    for line in contents:
        if default_label is None:
            name = re.match( r'[\s]*DEFAULT (.*)', line, re.M|re.I)
            if name:
                default_label = name.group(1)
                # Update the default label
                if default_label != label:
                    line = re.sub( r'[\s]*DEFAULT .*', 'DEFAULT %s' % label, line)
        else:
            name = re.match( r'[\s]*LABEL (.*)', line, re.M|re.I)
            if name:
                current_label = name.group(1)
                current_mlabel = None
                current_linux = None
                current_append = None

            name = re.match( r'[\s]*MENU LABEL (.*)', line, re.M|re.I)
            if name:
                current_mlabel = name.group(1)

            name = re.match( r'[\s]*LINUX (.*)', line, re.M|re.I)
            if name:
                current_linux = name.group(1)

            name = re.match( r'[\s]*APPEND (.*)', line, re.M|re.I)
            if name:
                current_append = name.group(1)
                # Inherit the boot arguments from the last known good boot entry
                if default_label == current_label:
                    default_append = current_append
                if label == current_label and mlabel == current_mlabel and image_path == current_linux:
                    new_entry = False
                    if add_kernel_params or rm_kernel_params:
                        modified_append = update_kernel_params(current_append, add_kernel_params, rm_kernel_params)
                        line = re.sub( r'([\s]*APPEND )(.*)', r'\g<1>' + modified_append, line)

        out.append(line)

    # If last line is not blank, add one.
    if not out[-1].isspace():
        out.append('\n')

    if new_entry:
        if default_append is None:
            print("Could not find APPEND from the current default boot entry %s" % default_label)
            sys.exit(1)
        if add_kernel_params or rm_kernel_params:
            append = update_kernel_params(default_append, add_kernel_params, rm_kernel_params)
        else:
            append = default_append
        print("Creating new entry in %s" % EXTLINUX)
        out.append('LABEL %s' % label)
        out.append('\n\tMENU LABEL %s' % mlabel)
        out.append('\n\tLINUX %s' % image_path)
        out.append('\n\tINITRD /boot/initrd')
        out.append('\n\tAPPEND %s' % append)
        out.append('\n')
    else:
        print("Using the existing boot entry \'%s\'" % label)

    fd, path = tempfile.mkstemp()
    with open(path, 'w') as fout:
        for line in out:
            fout.write(line)
    os.chmod(path, 0o644)

    backup = EXTLINUX + '.nv-update-extlinux-backup'
    shutil.copyfile(EXTLINUX, backup)
    shutil.move(path, EXTLINUX)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('variant', choices=['generic', 'real-time'], help='kernel variant')
    parser.add_argument('-a', '--add-param', action='append', metavar='param', dest='add_kernel_params', default=[],
                        help='kernel cmdline parameter to be appended')
    parser.add_argument('-r', '--rm-param', action='append', metavar='param', dest='rm_kernel_params', default=[],
                        help='kernel cmdline parameter to be removed')

    args = parser.parse_args()

    if args.variant == 'real-time':
        add_extlinux_entry('/boot/Image.real-time', 'real-time', 'real-time kernel', args.add_kernel_params, args.rm_kernel_params)
    else:
        add_extlinux_entry('/boot/Image', 'primary', 'primary kernel', args.add_kernel_params, args.rm_kernel_params)
