#!/usr/bin/env python3

# Copyright 2017 Red Hat, Inc.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import glob
import logging
import os
import re
import subprocess
import yaml


logging.basicConfig(
    level=logging.DEBUG if int(os.getenv('DIB_DEBUG_TRACE', 0)) > 0
    else logging.INFO)


def main():
    logging.info("Starting dracut regeneration")
    logging.debug("Debug logging enabled")

    dracut_env = os.getenv('DIB_DRACUT_ENABLED_MODULES')
    dracut_objects = yaml.safe_load(dracut_env)

    logging.debug(dracut_objects)

    modules_to_boot = []
    for dracut_object in dracut_objects:
        # first, install dependent packages
        packages = dracut_object.get('packages', [])
        for package in packages:
            cmdline = ["install-packages", package]
            logging.debug("Calling: %s" % cmdline)
            subp = subprocess.Popen(cmdline, stdout=subprocess.PIPE)
            out = subp.communicate()[0]
            if subp.returncode:
                e = subprocess.CalledProcessError(subp.returncode, cmdline)
                e.output = out
                raise e

        # second, compose the list of modules to boot
        modules_to_boot.append(dracut_object.get('name', None))

    conf_overrides = glob.glob('/etc/dracut.conf.d/*.conf')

    # regenerate dracut with the list of installed modules and conf overrides
    if modules_to_boot or conf_overrides:
        cmdline = ["select-boot-kernel-initrd"]
        logging.debug("Calling: %s" % cmdline)
        subp = subprocess.Popen(cmdline, stdout=subprocess.PIPE)
        out, err = subp.communicate()
        if subp.returncode:
            e = subprocess.CalledProcessError(subp.returncode, cmdline)
            e.output = out
            raise e

        kernel_set = out.decode().split(':')
        kernel_search = re.match("vmlinuz-(.*)", kernel_set[0])
        kernel_version = "%s" % kernel_search.groups(1)
        ramdisk_path = "/boot/%s" % kernel_set[1].strip()

        # list installed modules of dracut
        cmdline = ['dracut', '--list-modules']
        logging.debug("Calling: %s" % cmdline)
        subp = subprocess.Popen(cmdline, stdout=subprocess.PIPE)
        out, err = subp.communicate()
        if subp.returncode:
            e = subprocess.CalledProcessError(subp.returncode, cmdline)
            e.output = out
            raise e

        cmdline = ['dracut', '--force']

        if modules_to_boot:
            modules_to_boot = ' ' .join(modules_to_boot)
            cmdline += ['--add', modules_to_boot]
        cmdline += ['-f', ramdisk_path, kernel_version]
        logging.debug("Calling: %s" % cmdline)
        subp = subprocess.Popen(cmdline, stdout=subprocess.PIPE)
        out, err = subp.communicate()
        if subp.returncode:
            e = subprocess.CalledProcessError(subp.returncode, cmdline)
            e.output = out
            raise e

    logging.info("dracut regenerate done")


if __name__ == '__main__':
    main()

# Local Variables:
# mode: python
# End: