#!/usr/bin/env python3

# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# 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 os
import re
import shutil
import sys
import tempfile

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

def print_usage(script_name):
    name = os.path.basename(script_name)
    print("Usage: %s generic\n"
        "       %s real-time" % (name, name))

def add_extlinux_entry(image_path, label, mlabel):
    # 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

        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)
        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' % default_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__':

    if len(sys.argv) <= 1:
        print_usage(sys.argv[0])
        sys.exit(1)

    if sys.argv[1] == 'real-time':
        add_extlinux_entry('/boot/Image.real-time', 'real-time', 'real-time kernel')
    elif sys.argv[1] == 'generic':
        add_extlinux_entry('/boot/Image', 'primary', 'primary kernel')
    else:
        print_usage(sys.argv[0])
        sys.exit(1)
