#!/usr/bin/python3
#
# Copyright (C) 2019 Red Hat, Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
import argparse
import logging as log
import os
import shutil
import subprocess
import tempfile
import sys

from pylorax.mount import IsoMountpoint
from pylorax.treebuilder import udev_escape

# Constants for efimode (MACBOOT implies EFIBOOT)
NO_EFI = 0
EFIBOOT = 1
MACBOOT = 2

class Tool():
    """A class to check for executables and required files"""
    tools = []
    paths = {}
    requirements = []
    arches = []

    def __init__(self):
        # If there are arches listed it must be running on one of them
        if self.arches and os.uname().machine not in self.arches:
            log.debug("%s is not supported by this tool", os.uname().machine)
            raise RuntimeError("%s not supported" % os.uname().machine)

        # Check the system to see if the tools are available and record their paths
        for e in self.tools:
            self.paths[e] = shutil.which(e)
            if not self.paths[e]:
                log.debug("No executable %s found in PATH", e)
                raise RuntimeError("Missing executable")

        # Check to make sure all the required files are available
        for f in self.requirements:
            if not os.path.exists(f):
                log.debug("Missing file %s", f)
                raise RuntimeError("Missing requirement %s" % f)


class MkefibootTool(Tool):
    """Create the efiboot.img needed for EFI booting"""
    tools = ["mkefiboot"]

    def run(self, isodir, tmpdir, product="Fedora"):
        grafts = []
        # May have already been copied and added to grafts
        if not os.path.exists(os.path.join(tmpdir, "images")):
            # Copy over the whole images directory because efiboot.img is changing
            shutil.copytree(os.path.join(isodir, "images"), os.path.join(tmpdir, "images"))
            grafts = [(os.path.join(tmpdir, "images"), "images")]

        cmd = ["mkefiboot", "--label=ANACONDA"]
        if log.root.level < log.INFO:
            cmd.append("--debug")
        cmd.extend([os.path.join(tmpdir, "EFI/BOOT"), os.path.join(tmpdir, "images/efiboot.img")])
        log.debug(" ".join(cmd))
        try:
            subprocess.check_output(cmd)
        except subprocess.CalledProcessError as e:
            log.error(str(e))
            raise RuntimeError("Running mkefiboot")

        return grafts


class MkmacbootTool(MkefibootTool):
    """Create the macboot.img needed to EFI boot on Mac hardware"""
    tools = ["mkefiboot"]
    # NOTE: Order is important
    requirements = ["/usr/share/pixmaps/bootloader/fedora.icns",
                    "/usr/share/pixmaps/bootloader/fedora-media.vol"]

    def run(self, isodir, tmpdir, product="Fedora"):
        cmd = ["mkefiboot", "--label", "ANACONDA",
               "--apple", "--icon", self.requirements[0],
               "--diskname", self.requirements[1],
               "--product", product]
        if log.root.level < log.INFO:
            cmd.append("--debug")
        cmd.extend([os.path.join(tmpdir, "EFI/BOOT"), os.path.join(tmpdir, "images/macboot.img")])
        log.debug(" ".join(cmd))
        try:
            subprocess.check_output(cmd)
        except subprocess.CalledProcessError as e:
            log.error(str(e))
            raise RuntimeError("Running mkefiboot --apple")

        # images/macboot.img always exists with images/efiboot.img, already in grafts list
        return []


class MakeISOTool(Tool):
    """Class to hold details for specific iso creation tools"""
    def check_file_sizes(self, grafts):
        """Return True any file exceeds 4GiB"""
        for src, _ in grafts:
            if os.path.isdir(src):
                for top, dirs, files in os.walk(src):
                    for f in files + dirs:
                        if os.stat(os.path.join(top,f)).st_size >= 4*1024**3:
                            return True
            else:
                if os.stat(src).st_size >= 4*1024**3:
                    return True

        return False

    def _exec(self, cmd, grafts, output_iso, efimode=NO_EFI, implantmd5=True):
        """Add the grafts and run the command and then implant the md5 checksums"""
        cmd.append("-graft-points")

        for src, dest in grafts:
            cmd.append("%s=%s" % (dest, src))

        log.debug(" ".join(cmd))
        try:
            subprocess.check_output(cmd)
        except subprocess.CalledProcessError as e:
            log.error(str(e))
            raise RuntimeError("New ISO creation failed")

        if efimode > NO_EFI:
            cmd = ["isohybrid", "--uefi"]
            if efimode == MACBOOT:
                cmd.append("--mac")
            if log.root.level < log.INFO:
                cmd.append("--verbose")
            cmd.append(output_iso)

            log.debug(" ".join(cmd))
            try:
                subprocess.check_output(cmd)
            except subprocess.CalledProcessError as e:
                log.error(str(e))
                raise RuntimeError("Hybrid ISO creation failed")

        if not implantmd5:
            return

        cmd = ["implantisomd5", output_iso]
        log.debug(" ".join(cmd))
        try:
            subprocess.check_output(cmd)
        except subprocess.CalledProcessError as e:
            log.error(str(e))
            raise RuntimeError("implantisomd5 failed")


class Mkisofs_aarch64(MakeISOTool):
    """Use the mkisofs tool to create the final iso (aarch64)"""
    tools = ["mkisofs", "implantisomd5"]
    arches = ["aarch64", "arm"]

    def run(self, tmpdir, grafts, volume_name, output_iso, efimode=NO_EFI):
        cmd = ["mkisofs", "-o", output_iso, "-R", "-J", "-V", volume_name,
               "-T"]
        if log.root.level < log.INFO:
            cmd.append("--verbose")
        if efimode > NO_EFI:
            cmd.extend(["-eltorito-alt-boot", "-e", "images/efiboot.img", "-no-emul-boot"])
        if efimode > EFIBOOT:
            cmd.extend(["-eltorito-alt-boot", "-e", "images/macboot.img", "-no-emul-boot"])

        if self.check_file_sizes(grafts):
            cmd.append("-allow-limited-size")

        # Create the iso and implant the md5 checksums
        self._exec(cmd, grafts, output_iso, implantmd5=True)


class Mkisofs_ppc(MakeISOTool):
    """Use the mkisofs tool to create the final iso (ppc)"""
    tools = ["mkisofs"]
    requirements = ["/usr/share/lorax/templates.d/99-generic/config_files/ppc/mapping"]
    arches = ["ppc"]

    def run(self, tmpdir, grafts, volume_name, output_iso, efimode=NO_EFI):
        cmd = ["mkisofs", "-o", output_iso, "-R", "-J", "-V", volume_name,
               "-U", "-T", "-part", "-hfs", "-r", "-l", "-sysid", "PPC",
               "-chrp-boot", "-no-desktop", "-allow-multidot",
               "-map", self.requirements[0], "-hfs-bless", "boot/grub/powerpc-ieee1275"]
        if log.root.level < log.INFO:
            cmd.append("--verbose")

        if self.check_file_sizes(grafts):
            cmd.append("-allow-limited-size")

        # Create the iso and implant the md5 checksums
        self._exec(cmd, grafts, output_iso, efimode, implantmd5=True)


class Mkisofs_ppc64le(MakeISOTool):
    """Use the mkisofs tool to create the final iso (ppc64le)"""
    tools = ["mkisofs"]
    requirements = ["/usr/share/lorax/templates.d/99-generic/config_files/ppc/mapping"]
    arches = ["ppc64le"]

    def run(self, tmpdir, grafts, volume_name, output_iso, efimode=NO_EFI):
        cmd = ["mkisofs", "-o", output_iso, "-R", "-J", "-V", volume_name,
               "-T", "-part", "-hfs", "-r", "-l", "-sysid", "PPC",
               "-chrp-boot", "-no-desktop", "-allow-multidot",
               "-map", self.requirements[0]]
        if log.root.level < log.INFO:
            cmd.append("--verbose")

        if self.check_file_sizes(grafts):
            cmd.append("-allow-limited-size")

        # Create the iso and implant the md5 checksums
        self._exec(cmd, grafts, output_iso, efimode, implantmd5=True)


class Mkisofs_s390(MakeISOTool):
    """Use the mkisofs tool to create the final iso (s390)"""
    tools = ["mkisofs"]
    requirements = []
    arches = ["s390", "s390x"]

    def run(self, tmpdir, grafts, volume_name, output_iso, efimode=NO_EFI):
        cmd = ["mkisofs", "-o", output_iso, "-R", "-J", "-V", volume_name,
               "-b", "images/cdboot.img",
               "-c", "images/boot.cat",
               "-boot-load-size", "4", "-no-emul-boot"]
        if log.root.level < log.INFO:
            cmd.append("--verbose")

        if self.check_file_sizes(grafts):
            cmd.append("-allow-limited-size")

        # Create the iso and implant the md5 checksums
        self._exec(cmd, grafts, output_iso, efimode, implantmd5=True)


class Mkisofs_x86_64(MakeISOTool):
    """Use the mkisofs tool to create the final iso (x86_64)"""
    tools = ["mkisofs", "isohybrid"]
    requirements = []
    arches = ["x86_64", "i386"]

    def run(self, tmpdir, grafts, volume_name, output_iso, efimode=NO_EFI):
        cmd = ["mkisofs", "-o", output_iso, "-R", "-J", "-V", volume_name,
               "-b", "isolinux/isolinux.bin",
               "-c", "isolinux/boot.cat",
               "-boot-load-size", "4", "-boot-info-table", "-no-emul-boot"]
        if log.root.level < log.INFO:
            cmd.append("--verbose")
        if efimode > NO_EFI:
            cmd.extend(["-eltorito-alt-boot", "-e", "images/efiboot.img", "-no-emul-boot"])
        if efimode > EFIBOOT:
            cmd.extend(["-eltorito-alt-boot", "-e", "images/macboot.img", "-no-emul-boot"])

        if self.check_file_sizes(grafts):
            cmd.append("-allow-limited-size")

        # Create the iso, implant the md5 checksums, and create hybrid iso
        self._exec(cmd, grafts, output_iso, efimode, implantmd5=True)


class Xorrisofs_aarch64(MakeISOTool):
    """Use the xorrisofs tool to create the final iso (aarch64)"""
    tools = ["xorrisofs", "implantisomd5"]
    requirements = []
    arches = ["aarch64", "arm"]

    def run(self, tmpdir, grafts, volume_name, output_iso, efimode):
        cmd = ["xorrisofs", "-o", output_iso,
               "-R", "-J", "-V", volume_name]
        if log.root.level < log.INFO:
            cmd.append("--verbose")
        if efimode == EFIBOOT:
            cmd.extend(["-eltorito-alt-boot", "-e", "images/efiboot.img", "-no-emul-boot"])
        if efimode == MACBOOT:
            cmd.extend(["-eltorito-alt-boot", "-e", "images/macboot.img", "-no-emul-boot"])

        if self.check_file_sizes(grafts):
            cmd.extend(["-iso-level", "3"])

        # Create the iso and implant the md5 checksums
        self._exec(cmd, grafts, output_iso, implantmd5=True)


class Xorrisofs_ppc64le(MakeISOTool):
    """Use the xorrisofs tool to create the final iso (ppc64)"""
    tools = ["xorrisofs", "implantisomd5"]
    requirements = []
    arches = ["ppc64le"]

    def run(self, tmpdir, grafts, volume_name, output_iso, efimode=NO_EFI):
        cmd = ["xorrisofs", "-o", output_iso,
               "-R", "-J", "-V", volume_name,
               "-U", "-r", "-l", "-sysid", "PPC",
               "-A", volume_name, "-chrp-boot"]
        if log.root.level < log.INFO:
            cmd.append("--verbose")

        if self.check_file_sizes(grafts):
            cmd.extend(["-iso-level", "3"])

        # Create the iso and implant the md5 checksums
        self._exec(cmd, grafts, output_iso, efimode=NO_EFI, implantmd5=True)


class Xorrisofs_s390(MakeISOTool):
    """Use the xorrisofs tool to create the final iso (s390)"""
    tools = ["xorrisofs", "implantisomd5"]
    requirements = []
    arches = ["s390", "s390x"]

    def run(self, tmpdir, grafts, volume_name, output_iso, efimode=NO_EFI):
        cmd = ["xorrisofs", "-o", output_iso,
               "-R", "-J", "-V", volume_name,
               "-b", "images/cdboot.img", "-c", "images/boot.cat",
               "-boot-load-size", "4", "-no-emul-boot"]
        if log.root.level < log.INFO:
            cmd.append("--verbose")

        if self.check_file_sizes(grafts):
            cmd.extend(["-iso-level", "3"])

        # Create the iso and implant the md5 checksums
        self._exec(cmd, grafts, output_iso, efimode=NO_EFI, implantmd5=True)


class Xorrisofs_x86_64(MakeISOTool):
    """Use the xorrisofs tool to create the final iso"""
    tools = ["xorrisofs", "implantisomd5"]
    requirements = ["/usr/share/syslinux/isohdpfx.bin"]
    arches = ["x86_64", "i386"]

    def run(self, tmpdir, grafts, volume_name, output_iso, efimode=NO_EFI):
        cmd = ["xorrisofs", "-o", output_iso,
               "-R", "-J", "-V", volume_name,
               "-isohybrid-mbr", self.requirements[0],
               "-b", "isolinux/isolinux.bin",
               "-c", "isolinux/boot.cat",
               "-boot-load-size", "4", "-boot-info-table", "-no-emul-boot"]
        if log.root.level < log.INFO:
            cmd.append("--verbose")
        if efimode == EFIBOOT:
            cmd.extend(["-eltorito-alt-boot", "-e", "images/efiboot.img", "-no-emul-boot", "-isohybrid-gpt-basdat"])
        if efimode == MACBOOT:
            cmd.extend(["-eltorito-alt-boot", "-e", "images/macboot.img", "-no-emul-boot", "-isohybrid-gpt-hfsplus"])

        if self.check_file_sizes(grafts):
            cmd.extend(["-iso-level", "3"])

        # Create the iso and implant the md5 checksums
        self._exec(cmd, grafts, output_iso, efimode=NO_EFI, implantmd5=True)


# xorrisofs based classes
XorrisofsTools = [Xorrisofs_aarch64, Xorrisofs_ppc64le, Xorrisofs_s390, Xorrisofs_x86_64]

# mkisofs based classes
MkisofsTools = [Mkisofs_aarch64, Mkisofs_ppc, Mkisofs_ppc64le, Mkisofs_s390, Mkisofs_x86_64]


class MakeKickstartISO():
    def __init__(self, ks, input_iso, output_iso, add_paths, cmdline=""):
        self.ks = ks
        self.input_iso = input_iso
        self.output_iso = output_iso
        self.add_paths = add_paths
        self.isotool = None
        self._cmdline = cmdline
        self.label = ""
        self.iso = None
        self.mkmacboot = None
        self.efimode = NO_EFI

        errors = False
        for f in [ks, input_iso] + add_paths:
            if not os.path.exists(f):
                log.error("%s is missing", f)
                errors = True

        if errors:
            raise RuntimeError("Missing Files")

        # Mount the iso so we can gather information from it
        try:
            self.iso = IsoMountpoint(self.input_iso)
            self.label = self.iso.label
            log.info("Volume Id = %s", self.label)

            # Get the list of files in / of the source iso, and their full paths
            iso_files = os.listdir(self.iso.mount_dir)
            log.info("ISO files = %s", iso_files)

            # Populate the grafts with the existing ISO, except EFI, images, and isolinux
            # which are copied to tmpdir for editing
            skip_iso = ["EFI", "images", "isolinux", "boot.cat", "boot.catalog", "TRANS.TBL"]
            self.grafts = [(os.path.join(self.iso.mount_dir, f), f) for f in iso_files if f not in skip_iso]

            if os.path.exists(os.path.join(self.iso.mount_dir, "images/efiboot.img")):
                self.efimode = EFIBOOT
                self.mkefiboot = MkefibootTool()
            if os.path.exists(os.path.join(self.iso.mount_dir, "images/macboot.img")):
                self.efimode = MACBOOT
                self.mkmacboot = MkmacbootTool()

            for t in XorrisofsTools + MkisofsTools:
                try:
                    self.isotool = t()
                    break
                except RuntimeError:
                    continue
            if not self.isotool:
                raise RuntimeError("No suitable iso creation tool found.")

            log.info("Using %s to create the new iso", self.isotool.tools[0])
        except:
            if self.iso:
                self.iso.umount()
            raise

    @property
    def escaped_label(self):
        """Return the escaped iso volume label"""
        return udev_escape(self.label)

    def close(self):
        if self.iso:
            self.iso.umount()

    @property
    def add_args(self,):
        """Return the arguments to add to the config file kernel cmdline"""
        return "inst.ks=hd:LABEL=%s:/%s %s" % (self.escaped_label, os.path.basename(self.ks), self._cmdline)

    def remove_generated_files(self, grafts):
        """Remove the files generated by the iso creation tools"""
        catalog_files = ["boot.cat", "boot.catalog", "TRANS.TBL"]
        for src, _ in grafts:
            if not os.path.isdir(src):
                continue
            for f in catalog_files:
                if os.path.exists(os.path.join(src, f)):
                    os.unlink(os.path.join(src, f))

    def run_mkefiboot(self, isodir, tmpdir):
        grafts = self.mkefiboot.run(isodir, tmpdir)
        if self.efimode == MACBOOT:
            grafts.extend(self.mkmacboot.run(isodir, tmpdir))

        return grafts

    def edit_configs(self, isodir, tmpdir):
        """Find and edit any configuration files

        Add the inst.ks= argument plus extra cmdline arguments
        """
        # Note that some of these may not exist, depending on the arch being used
        grafts = self._edit_isolinux(isodir, tmpdir)
        grafts.extend(self._edit_efi(isodir, tmpdir))
        grafts.extend(self._edit_ppc(isodir, tmpdir))
        grafts.extend(self._edit_s390(isodir, tmpdir))

        return grafts

    def _edit_isolinux(self, isodir, tmpdir):
        """Copy the isolinux.cfg file and add the cmdline args"""

        orig_cfg = os.path.join(isodir, "isolinux/isolinux.cfg")
        if not os.path.exists(orig_cfg):
            log.warning("No isolinux/isolinux.cfg file found")
            return []

        # Copy over the whole directory
        shutil.copytree(os.path.join(isodir, "isolinux"), os.path.join(tmpdir, "isolinux"))

        # Edit the config file
        with open(orig_cfg, "r") as in_fp:
            with open(os.path.join(tmpdir, "isolinux/isolinux.cfg"), "w") as out_fp:
                for line in in_fp:
                    out_fp.write(line.rstrip("\n"))
                    if "append" in line:
                        out_fp.write(" "+self.add_args)
                    out_fp.write("\n")

        # Return the graft src, dest
        return [(os.path.join(tmpdir, "isolinux"), "isolinux")]

    def _edit_efi(self, isodir, tmpdir):
        """Copy the efi config files and add the cmdline args"""
        # At least one of these must be present
        efi_cfgs = ["EFI/BOOT/grub.cfg", "EFI/BOOT/BOOT.conf"]

        if not os.path.exists(os.path.join(isodir, "EFI")):
            log.warning("No EFI directory file found")
            return []

        # mkefiboot needs to use the EFI/ directory along with the modified config files
        # so start by copying the whole directory over
        shutil.copytree(os.path.join(isodir, "EFI"), os.path.join(tmpdir, "EFI"))

        found_cfg = False
        for cfg in efi_cfgs:
            orig_cfg = os.path.join(isodir, cfg)
            if not os.path.exists(orig_cfg):
                continue

            dest_cfg = os.path.join(tmpdir, cfg)
            with open(orig_cfg, "r") as in_fp:
                with open(dest_cfg, "w") as out_fp:
                    for line in in_fp:
                        out_fp.write(line.rstrip("\n"))
                        # Some start with linux (aarch64), others with linuxefi (x86_64)
                        if line.strip().startswith("linux"):
                            out_fp.write(" "+self.add_args)
                        out_fp.write("\n")
            found_cfg = True

        if not found_cfg:
            raise RuntimeError("ISO is missing the EFI config files")

        # Graft the copy of the EFI directory
        return [(os.path.join(tmpdir, "EFI"), "EFI")]

    def _edit_ppc(self, isodir, tmpdir):
        """Edit the boot/grub/grub.cfg file, adding the kickstart and extra arguments"""
        orig_cfg = os.path.join(isodir, "boot/grub/grub.cfg")
        if not os.path.exists(orig_cfg):
            log.warning("No boot/grub/grub.cfg file found")
            return []

        # Copy over the whole directory
        shutil.copytree(os.path.join(isodir, "boot"), os.path.join(tmpdir, "boot"))

        # Edit the config file
        with open(orig_cfg, "r") as in_fp:
            with open(os.path.join(tmpdir, "boot/grub/grub.cfg"), "w") as out_fp:
                for line in in_fp:
                    out_fp.write(line.rstrip("\n"))
                    if line.strip().startswith("linux "):
                        out_fp.write(" "+self.add_args)
                    out_fp.write("\n")

        # Graft the copy of the boot directory
        return [(os.path.join(tmpdir, "boot"), "boot")]

    def _edit_s390(self, isodir, tmpdir):
        """Edit the images/generic.prm file, adding the kickstart and extra arguments"""
        orig_cfg = os.path.join(isodir, "images/generic.prm")
        if not os.path.exists(orig_cfg):
            log.warning("No images/generic.prm file found")
            return []

        # Copy over the whole directory
        shutil.copytree(os.path.join(isodir, "images"), os.path.join(tmpdir, "images"))

        # Append to the config file
        with open(os.path.join(tmpdir, "images/generic.prm"), "a") as out_fp:
            out_fp.write(self.add_args+"\n")

        # Graft the copy of the images directory
        return [(os.path.join(tmpdir, "images"), "images")]

    def run(self):
        """Modify the ISO"""
        try:
            # Make a temporary directory to hold modified files
            with tempfile.TemporaryDirectory(prefix="mkksiso-") as tmpdir:

                # Copy and edit the configuration files
                self.grafts.extend(self.edit_configs(self.iso.mount_dir, tmpdir))

                # Run the mkefiboot tool on the edited EFI directory, add the new files to the grafts
                self.grafts.extend(self.run_mkefiboot(self.iso.mount_dir, tmpdir))

                # Add the kickstart to grafts
                self.grafts.extend([(self.ks, os.path.basename(self.ks))])

                # Add the extra files to grafts
                self.grafts.extend([(f, os.path.basename(f)) for f in self.add_paths])

                # Remove files that will be regenerated by the iso creation process
                self.remove_generated_files(self.grafts)

                log.info("grafts = %s", self.grafts)
                self.isotool.run(tmpdir, self.grafts, self.label, self.output_iso, self.efimode)
        finally:
            self.close()


def setup_args():
    """ Return argparse.Parser object of cmdline."""
    parser = argparse.ArgumentParser(description="Add a kickstart and files to an iso")

    parser.add_argument("-a", "--add", action="append", dest="add_paths", default=[],
                        type=os.path.abspath,
                        help="File or directory to add to ISO (may be used multiple times)")
    parser.add_argument("-c", "--cmdline", dest="cmdline", metavar="CMDLINE",
                        help="Arguments to add to kernel cmdline")
    parser.add_argument("--debug", action="store_const", const=log.DEBUG,
                        dest="loglevel", default=log.INFO,
                        help="print debugging info")

    parser.add_argument("ks", type=os.path.abspath, help="Kickstart to add to the ISO")
    parser.add_argument("input_iso", type=os.path.abspath, help="ISO to modify")
    parser.add_argument("output_iso", type=os.path.abspath, help="Full pathname of iso to be created")
    args = parser.parse_args()

    return args


def main():
    args = setup_args()
    log.basicConfig(format='%(levelname)s:%(message)s', level=args.loglevel)

    if os.getuid() != 0:
        log.error("You must run this as root, it needs to mount the iso and run mkefiboot")
        sys.exit(-1)

    try:
        app = MakeKickstartISO(args.ks, args.input_iso, args.output_iso,
                               args.add_paths, args.cmdline)
        app.run()
    except RuntimeError as e:
        log.error(str(e))

if __name__ == '__main__':
    main()