#!/usr/bin/python
#
# iso-to-pxe: Extract the files needed to PXE boot from an iso
#
# Copyright (C) 2012  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/>.
#
# Author(s): Brian C. Lane <bcl@redhat.com>
#
import os
import sys
import argparse
import shutil
from pylorax.sysutils import joinpaths
from pylorax.executils import execWithRedirect, execWithCapture
import tempfile

# ==== SETUP ====
# Full path to top level of the PXE server's image directory
IMAGE_DIR = "/extra/redhat/tftpboot/images/"
# Relative path to use in the PXE config
PXE_TOP_DIR = "images/"
# URL pointing to IMAGE_DIR, for retrieving LiveOS/squashfs.img
IMAGE_URL = "http://proxy.brianlane.com/"

# Directory of PXE config files, named DEF_DIST
PXE_CFG_DIR = "/extra/redhat/tftpboot/pxelinux.cfg/"

# This is appended to IMAGE_DIR as a path and to PXE_CFG_DIR as a config file
DEF_DIST = "fedora"
# This is appended to DEF_DIST and can be a path
DEF_RELEASE = "18"

# Default extra args to add to the APPEND line
DEF_ARGS = ""


# from livemedia-creator
class IsoMountpoint(object):
    """
    Mount the iso on a temporary directory and check to make sure the
    vmlinuz and initrd.img files exist
    Check the iso for a LiveOS directory and set a flag.
    Extract the iso's label.
    """
    def __init__( self, iso_path ):
        self.label = None
        self.iso_path = iso_path
        self.mount_dir = tempfile.mkdtemp()
        if not self.mount_dir:
            raise Exception("Error creating temporary directory")

        cmd = ["mount", "-o", "loop", self.iso_path, self.mount_dir]
        print(cmd)
        execWithRedirect( cmd[0], cmd[1:] )

        self.kernel = self.mount_dir+"/isolinux/vmlinuz"
        self.initrd = self.mount_dir+"/isolinux/initrd.img"

        if os.path.isdir( self.mount_dir+"/repodata" ):
            self.repo = self.mount_dir
        else:
            self.repo = None
        self.liveos = os.path.isdir( self.mount_dir+"/LiveOS" )

        try:
            for f in [self.kernel, self.initrd]:
                if not os.path.isfile(f):
                    raise Exception("Missing file on iso: {0}".format(f))
        except:
            self.umount()
            raise

#        self.get_iso_label()

    def umount( self ):
        cmd = ["umount", self.mount_dir]
        print(cmd)
        execWithRedirect( cmd[0], cmd[1:] )
        os.rmdir( self.mount_dir )

    def get_iso_label( self ):
        """
        Get the iso's label using isoinfo
        """
        cmd = [ "isoinfo", "-d", "-i", self.iso_path ]
        print(cmd)
        isoinfo_output = execWithCapture( cmd[0], cmd[1:] )
        print(isoinfo_output)
        for line in isoinfo_output.splitlines():
            if line.startswith("Volume id: "):
                self.label = line[11:]
                return


def make_parents(path):
    """ Make parent directories
    """
    if not os.path.isdir(path):
        os.makedirs(path)


def run():
    """ run iso-to-pxe
    """
    parser = argparse.ArgumentParser(description="Setup PXE from .iso")
    parser.add_argument("-d", "--imagedir", default=IMAGE_DIR,
                        help="Full path to top level image directory")
    parser.add_argument("-t", "--pxe-top", default=PXE_TOP_DIR,
                        help="PXE top level directory")
    parser.add_argument("-u", "--url", default=IMAGE_URL,
                        help="URL for stage2")
    parser.add_argument("-c", "--config", default=PXE_CFG_DIR,
                        help="PXE config file directory")
    parser.add_argument("-D", "--dist", default=DEF_DIST,
                        help="Distribution directory/config name")
    parser.add_argument("-r", "--releasever", default=DEF_RELEASE,
                        help="Release version directory")
    parser.add_argument("-s", "--skip-pxe", default=False,
                        help="Skip appending PXE config")
    parser.add_argument("-x", "--extra", default=DEF_ARGS,
                        help="Extra args for PXE APPEND line")
    parser.add_argument("iso", type=os.path.abspath,
                        help="path to iso")
    args = parser.parse_args()

    if os.geteuid() != 0:
        print("You must run this as root")
        sys.exit(1)

    # mount iso on a temp dir
    iso = IsoMountpoint(args.iso)

    # copy the files
    files = []
    initrd_src = joinpaths(iso.mount_dir, "isolinux/initrd.img")
    initrd_dest = joinpaths(args.imagedir, args.dist, args.releasever)
    files.append((initrd_src, initrd_dest))

    vmlinuz_src = joinpaths(iso.mount_dir, "isolinux/vmlinuz")
    vmlinuz_dest = joinpaths(args.imagedir, args.dist, args.releasever)
    files.append((vmlinuz_src, vmlinuz_dest))

    liveos_src = joinpaths(iso.mount_dir, "LiveOS/squashfs.img")
    liveos_dest = joinpaths(args.imagedir, args.dist, args.releasever, "LiveOS")
    files.append((liveos_src, liveos_dest))

    for src, dest in files:
        make_parents(dest)
        print("Copying %s to %s" % (src, dest))
        shutil.copy2(src, dest)

    if not args.skip_pxe:
        pxe_config = """
LABEL %(dist)s %(releasever)s
    MENU  %(dist)s %(releasever)s
    KERNEL %(pxe_top)s/%(dist)s/%(releasever)s/vmlinuz
    APPEND initrd=%(pxe_top)s/%(dist)s/%(releasever)s/initrd.img stage2=%(url)s/%(dist)s/%(releasever)s/ %(extra)s
""" % (vars(args))
        print pxe_config
        with open(joinpaths(args.config,args.dist), "a") as f:
            f.write(pxe_config)

    iso.umount()


if __name__ == "__main__":
    run()