From c15349c3da253a1464e6b81db1b7dd4861c00f5c Mon Sep 17 00:00:00 2001
From: "Brian C. Lane" <bcl@redhat.com>
Date: Thu, 11 Feb 2016 10:50:46 -0800
Subject: [PATCH] livemedia-creator: Add /usr/share/lorax/templates.d/ support

---
 src/pylorax/__init__.py    | 30 ++++++++++++++++++++++--------
 src/sbin/livemedia-creator |  9 ++++++---
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/src/pylorax/__init__.py b/src/pylorax/__init__.py
index e19264e3..a30eba30 100644
--- a/src/pylorax/__init__.py
+++ b/src/pylorax/__init__.py
@@ -148,14 +148,8 @@ class Lorax(BaseLoraxClass):
         Otherwise use the sharedir
         """
         if not self._templatedir:
-            templatedir = self.conf.get("lorax", "sharedir")
-            if os.path.isdir(joinpaths(templatedir, "templates.d")):
-                try:
-                    templatedir = sorted(glob(joinpaths(templatedir, "templates.d", "*")))[0]
-                except IndexError:
-                    pass
-            logger.info("Using templatedir %s", templatedir)
-            self._templatedir = templatedir
+            self._templatedir = find_templates(self.conf.get("lorax", "sharedir"))
+            logger.info("Using templatedir %s", self._templatedir)
         return self._templatedir
 
     def init_stream_logging(self):
@@ -424,3 +418,23 @@ def setup_logging(logfile, theLogger):
     fh = logging.FileHandler(filename=f, mode="w")
     fh.setLevel(logging.DEBUG)
     program_log.addHandler(fh)
+
+
+def find_templates(templatedir="/usr/share/lorax"):
+    """ Find the templates to use.
+
+    :param str templatedir: Top directory to search for templates
+    :returns: Path to templates
+    :rtype: str
+
+    If there is a templates.d directory under templatedir the
+    lowest numbered directory entry is returned.
+
+    eg. /usr/share/lorax/templates.d/99-generic/
+    """
+    if os.path.isdir(joinpaths(templatedir, "templates.d")):
+        try:
+            templatedir = sorted(glob(joinpaths(templatedir, "templates.d", "*")))[0]
+        except IndexError:
+            pass
+    return templatedir
diff --git a/src/sbin/livemedia-creator b/src/sbin/livemedia-creator
index b7cd36d2..f6bf8837 100755
--- a/src/sbin/livemedia-creator
+++ b/src/sbin/livemedia-creator
@@ -45,7 +45,7 @@ from mako.template import Template
 from mako.exceptions import text_error_template
 
 # Use the Lorax treebuilder branch for iso creation
-from pylorax import ArchData, setup_logging
+from pylorax import ArchData, setup_logging, find_templates
 from pylorax.base import DataHolder
 from pylorax.treebuilder import TreeBuilder, RuntimeBuilder, udev_escape
 from pylorax.treebuilder import findkernels
@@ -1091,7 +1091,7 @@ def main():
                         type=os.path.abspath,
                         help="Name and path for primary logfile, other logs will "
                              "be created in the same directory.")
-    parser.add_argument("--lorax-templates", default="/usr/share/lorax/",
+    parser.add_argument("--lorax-templates", default=None,
                         type=os.path.abspath,
                         help="Path to mako templates for lorax")
     parser.add_argument("--tmp", default="/var/tmp", type=os.path.abspath,
@@ -1205,6 +1205,9 @@ def main():
 
     log.debug( opts )
 
+    # Find the lorax templates
+    opts.lorax_templates = find_templates(opts.lorax_templates or "/usr/share/lorax")
+
     # Check for invalid combinations of options, print all the errors and exit.
     errors = []
     if not opts.disk_image and not opts.fs_image and not opts.ks:
@@ -1213,7 +1216,7 @@ def main():
     if opts.ks and not os.path.exists(opts.ks[0]):
         errors.append("kickstart file (%s) is missing." % opts.ks[0])
 
-    if opts.make_iso and not os.path.exists( opts.lorax_templates):
+    if opts.make_iso and not os.path.exists(opts.lorax_templates):
         errors.append("The lorax templates directory (%s) doesn't "
                       "exist." % opts.lorax_templates)