#!/usr/bin/env python
#
# Author: Chris Jones
# Copyright 2017, LambdaStack
#
# 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.
#

# This script does everything from the beginning. It can be split apart to do only the package
# and kickstart addition and iso build by using a base directory of an already rsync directory.

import argparse
import datetime
import subprocess
import fnmatch
import os
import glob

# Set mount point of orginal iso. Change the directory if desired.
ISOBASEIMG = "rhel-server-7.2-x86_64-dvd.iso"
ISOFILE = "rhel-server-7.2-x86_64"
ISOEXT = ".iso"

# NOTE: This is VERY important starting in RHEL 7. It MUST be same as hd:LABEL= for it to work properly.
ISOVOLUMEID = "RHEL-7.2 Server.x86_64"

ISOBUILDPACKAGESDIR = "/tmp/build/packages"
ISOPOSTINSTALL = "postinstall"
ISOBASEDIR = "/tmp"
ISOBASEBUILDDIR = "/tmp/build"
ISOBUILDDIR = "/tmp/build/iso"
ISOKICKSTARTDIR = "/tmp/build/iso/ks"
#ISOBUILDTESTDIR = "/tmp/build/isotest"
ISOMNTPNT = "/tmp/build/mnt"

PACKAGER = "LambdaStack"
COMPS_FILE = "comps.xml"

# NOTE: Each RHEL build has a specific file used by Anaconda to index the Packages. Could use *-x86_64-comps.xml
# sha256sum of comps.xml generates the checksum of the file.
# Example file: 175ddec2056ec6b5ef267cea35f8ec679314afbfb019957e53f71725bcc5d829-c7-x86_64-comps.xml
# Default:
ANACONDAFILE = ""

# Since there are no interactive pieces, put the commands in an array.

# NOTE: sudo is required on some commands and optional on others below. If you move the ISOMNTPNT outside of user home
# directory then add sudo and anywhere else optional sudo commands are used if outside of home directory.
# rsync is used below but cp -av could also be used. The point is to make SURE .discinfo and .treeinfo are copied to
# the ISOBUILDDIR

# NOTE: Could have done mkdir -p ~/build/{iso,mnt,iso/ks} instead of all mkdir commands but it is easier to print out
# cmds for debugging this way.
# If we keep a 'base' directory used to regen iso files then this only has to be done once.
# Make sure .treeinfo and .discinfo is present - below we just cp them again to make sure.
# IMPORTANT directories in the build directory:
#  LiveOS, Packages, isolinux, images (and sub-directories), repodata (generated by repo_cmds below). The reason for
# these directories is because .treeinfo references them for stage2 loading.

init_cmds = [
        "sudo yum -y install anaconda anaconda-help anaconda-runtime",
        "sudo yum -y install createrepo /usr/bin/yumdownloader /usr/bin/mkisofs",
        "rm -rf %s" % ISOBASEBUILDDIR,
        "mkdir -p {%s,%s,%s}" % (ISOBASEDIR,ISOBUILDDIR,ISOMNTPNT),
        "sudo mount -o loop %s/%s %s" % (ISOBASEDIR, ISOBASEIMG, ISOMNTPNT),
        "rm -rf %s" % ISOBUILDDIR,
        "rsync -av %s/ %s/" % (ISOMNTPNT, ISOBUILDDIR),
        "rm -rf %s/repodata" % ISOBUILDDIR,
        "cp -av %s/.discinfo %s/.discinfo" % (ISOMNTPNT, ISOBUILDDIR),
        "cp -av %s/.treeinfo %s/.treeinfo" % (ISOMNTPNT, ISOBUILDDIR),
        "mkdir -p {%s,%s}" % (ISOKICKSTARTDIR, ISOPOSTINSTALL),
        "cp %s/repodata/%s %s/%s" % (ISOMNTPNT, ANACONDAFILE, ISOBUILDDIR, COMPS_FILE),
        "sudo umount %s" % ISOMNTPNT,
        "find %s -name TRANS.TBL -exec rm -f {} \;" % ISOBUILDDIR]
# NOTE: In the above 'comps.xml' you can use the exact one as the original iso - just find the first '-' and
# substring the name from the index.

rpm_cmds = []

kickstart_cmds = []

# If we keep a 'base' directory used to regen iso files then this only has to be done once. However, everytime we
# add/remove packages then this will have to be called so as to keep the package indexes updated.

repo_cmds = [
        "rm -rf %s/repodata" % ISOBUILDDIR,
        "discinfo=$(head -1 %s/.discinfo)" % ISOBUILDDIR,
        "createrepo -u \"media://$discinfo\" -g %s %s" % (COMPS_FILE, ISOBUILDDIR)]

# Commented out from iso_cmds below...
#  "chmod 664 %s/isolinux/*" % ISOBUILDDIR,
#  "chmod 664 %s/*" % ISOKICKSTARTDIR,

iso_cmds = [
        "rm %s/%s" % (ISOBUILDDIR, COMPS_FILE),
        "rm -f %s/%s%s" % (ISOBASEDIR, ISOFILE, ISOEXT),
        "mkisofs -r -J -N -d -hide-rr-moved -R -T -no-emul-boot -boot-load-size 4 -boot-info-table -V \"%s\" -p \"%s\" -A \"%s - %s\" -b isolinux/isolinux.bin -c isolinux/boot.cat -x \"lost+found\" -eltorito-alt-boot -e images/efiboot.img -no-emul-boot -o %s/%s%s %s" % (ISOVOLUMEID, PACKAGER, ISOVOLUMEID, datetime.date.today(), ISOBASEDIR, ISOFILE, ISOEXT, ISOBUILDDIR),
        "implantisomd5 %s/%s%s" % (ISOBASEDIR, ISOFILE, ISOEXT)]


def find_file(path, pattern):
    for file in os.listdir(path):
        if fnmatch.fnmatch(file, pattern):
            return file
    return None


def find_comps_file_name(file_name):
    if file_name is None:
        return None
    index = file_name.index('-')
    if index:
        return file_name[index-1:]
    return file_name

def run_commands(commands=None):
    for cmd in commands:
        subprocess.call(cmd, shell=True)


def init():
    if glob.glob("%s/.discinfo" % ISOBUILDDIR):
        subprocess.call("sudo umount %s" % ISOMNTPNT, shell=True)

    # Set ANACONDAFILE
    ANACONDAFILE=find_file('%s/repodata/' % ISOMNTPNT, '*comps*.xml')
    COMPS_FILE=find_comps_file_name(ANACONDAFILE)

    run_commands(init_cmds)


def add_rpms():
    # NOTE: Add/remove packages in the /Packages directory
    pass


def create_repo():
    run_commands(repo_cmds)


def add_kickstart():
    # NOTE: Add kickstart files to the /ks directory
    # IMPORTANT: Make sure you change the isolinux/isolinux.cfg to add/change menu item if there is a menu. Regardless,
    # put kickstart command after other 'append' items. For example:
    # append initrd=initrd.img inst.stage2=hd:LABEL=CentOS\x207\x20x86_64 ks=cdrom:/ks/rhel.ks
    # Also make sure ks=cdrom:/ks/rhel.ks (set directory to kickstart file). You can have one for each menu item.
    pass


def build_iso():
    run_commands(iso_cmds)


def main():
    p = argparse.ArgumentParser(description='ISO Builder', prog='bootstrap-iso-build')
    p.add_argument('--init', '-i', action='store_true', help='Initialize the build directory and move ISO mount data to the working directory.')
    p.add_argument('--addrpms', '-a', action='store_true', help='Copies/Creates RPM Packages to working directory.')
    p.add_argument('--createrepo', '-c', action='store_true', help='Creates Packages index in repodata directory.')
    p.add_argument('--addkickstart', '-k', action='store_true', help='Copies/Creates kickstart files in working directory.')
    p.add_argument('--buildiso', '-b', action='store_true', help='Builds ISO.')

    options = p.parse_args()

    if options.init:
        init()
    if options.addrpms:
        add_rpms()
    if options.createrepo:
        create_repo()
    if options.addkickstart:
        add_kickstart()
    if options.buildiso:
        build_iso()

    if not options.init and not options.addrpms and not options.createrepo and not options.addkickstart and not options.buildiso:
        p.print_help()

if __name__ == "__main__":
    main()
