#!/usr/bin/python3 -u

# SPDX-FileCopyrightText: 2023 - 2023 ENCRYPTED SUPPORT LP <adrelanos@kicksecure.com>
#
# SPDX-License-Identifier: AGPL-3.0-or-later

import sys
import re

def process_grub_line(input_line, extra_params, in_custom_section):
    leading_whitespace = input_line[:len(input_line) - len(input_line.lstrip())]

    if input_line.strip().startswith('menuentry'):
        pattern = re.compile(r"'gnulinux-(.+?)'")
        input_line = re.sub(pattern, "'gnulinux-iso'", input_line)

    if 'search --no-floppy --fs-uuid --set=root' in input_line:
        return leading_whitespace + '# ' + input_line.lstrip(), in_custom_section

    if in_custom_section:
        if input_line.strip() == "":
            return input_line, in_custom_section
        return leading_whitespace + '# ' + input_line.lstrip(), in_custom_section

    line_starts = ('linux', 'module2')
    for line_start in line_starts:
        if input_line.lstrip().startswith(line_start):
            parts = input_line.split(line_start)
            pre_command_whitespace = parts[0]
            post_command_whitespace = parts[1][:len(parts[1]) - len(parts[1].lstrip())]

            remaining_parts = parts[1].lstrip().split()

            excluded_startswith = ['root=UUID']
            ## Was functional:
            #excluded_exact = ['ro', 'quiet', 'rhgb']
            excluded_exact = ['ro']

            filtered_parts = []
            removed_parts = []
            for part in remaining_parts:
                if any(part.startswith(excluded) for excluded in excluded_startswith) or part in excluded_exact:
                    removed_parts.append(part)
                else:
                    filtered_parts.append(part)

            if removed_parts:
                comment = leading_whitespace + "# Removed options: " + ' '.join(removed_parts) + '\n'
                processed_line = comment + pre_command_whitespace + line_start + post_command_whitespace + ' '.join(filtered_parts) + ' ' + ' '.join(extra_params) + '\n'
                return processed_line, in_custom_section

    return input_line, in_custom_section  # Return the line unmodified if it doesn't need processing

def process_grub_config(input_file, output_file, extra_params):
    in_custom_section = False
    try:
        with open(input_file, 'r') as infile, open(output_file, 'w') as outfile:
            # Header comment
            header_comment = "# This file was auto generated by\n# grub-raw-cfg-to-grub-iso-cfg\n# based on:\n# /boot/grub/grub-raw-image-temp.cfg\n\n"
            outfile.write(header_comment)

            # Added options comment
            if extra_params:
                added_options_comment = "# Added options: " + ' '.join(extra_params) + "\n\n"
                outfile.write(added_options_comment)

            for line in infile:
                if "### BEGIN /etc/grub.d/11_linux_live ###" in line.strip():
                    outfile.write("# ISO is always Live. Remove superfluous grub-live boot entry.\n")
                    in_custom_section = True

                processed_line, in_custom_section = process_grub_line(line, extra_params, in_custom_section)
                if processed_line is not None:
                    outfile.write(processed_line)

                if in_custom_section and "### END /etc/grub.d/11_linux_live ###" in line.strip():
                    in_custom_section = False

    except IOError as e:
        print(f"Error: {e}")

def main():
    if len(sys.argv) < 4:
        print("Usage: grub-raw-cfg-to-grub-iso-cfg <input_grub.cfg> <output_grub-iso.cfg> <extra_param1> [<extra_param2> ...]")
        sys.exit(1)

    input_grub_file = sys.argv[1]
    output_grub_file = sys.argv[2]
    extra_kernel_params = sys.argv[3:]

    process_grub_config(input_grub_file, output_grub_file, extra_kernel_params)

if __name__ == "__main__":
    main()
