From 30e901ef19498b683c6d014a69e4b076cde1f956 Mon Sep 17 00:00:00 2001
From: Jarrod Johnson <jjohnson2@lenovo.com>
Date: Wed, 18 Aug 2021 17:34:27 -0400
Subject: [PATCH] Attempt -p and -r arguments to add repos and speciy alternate
 pakagke list.

---
 imgutil/imgutil | 33 ++++++++++++++++++++++++++-------
 1 file changed, 26 insertions(+), 7 deletions(-)

diff --git a/imgutil/imgutil b/imgutil/imgutil
index 8ce45b3b..dc39b0e2 100644
--- a/imgutil/imgutil
+++ b/imgutil/imgutil
@@ -324,7 +324,7 @@ def encrypt_image(plainfile, cryptfile, keyfile):
 
 
 
-def create_yumconf(sourcedir):
+def create_yumconf(sourcedir, addrepos):
     repodir = tempfile.mkdtemp(prefix='genimage-yumrepos.d-')
     yumconf = open(os.path.join(repodir, 'repos.repo'), 'w+')
     if '/' not in sourcedir:
@@ -346,6 +346,13 @@ def create_yumconf(sourcedir):
                 currdir = os.path.join(sourcedir, repopath)
                 yumconf.write('baseurl={0}\n'.format(currdir))
                 yumconf.write('enabled=1\ngpgcheck=0\n\n')
+    addrepoidx = 1
+    for repo in addrepos.split(','):
+        yumconf.write('[addrepo-{0}]\n', addrepoidx)
+        yumconf.write('name=Add-on repository {0}\n'.format(addrepoidx))
+        yumconf.write('baseurl={0}\n'.format(repo))
+        yumconf.write('enabled=1\ngpgcheck\0\n\n')
+        addrepoidx += 1
     return repodir
 
 def get_mydir(oscategory):
@@ -357,12 +364,19 @@ def get_mydir(oscategory):
     return gencopy
 
 class OsHandler(object):
-    def __init__(self, name, version, arch):
+    def __init__(self, name, version, arch, args):
         self.name = name
         self.version = version
         self.arch = arch
         self.sourcepath = None
         self.osname = '{}-{}-{}'.format(name, version, arch)
+        if args.pkglist:
+            self.pkglist = args.pkglist
+            if '/' not in self.pkglist:
+                self.pkglist = os.path.join(get_mydir(self.oscategory), self.pkglist)
+        else:
+            self.pkglist = os.path.join(get_mydir(self.oscategory), 'pkglist')
+        self.addrepos = args.addrepos
 
     def get_json(self):
         odata = [self.oscategory, self.version, self.arch, self.name]
@@ -374,7 +388,7 @@ class OsHandler(object):
         return json.dumps(info)
 
     def list_packages(self):
-        with open(os.path.join(get_mydir(self.oscategory), 'pkglist'), 'r') as pkglist:
+        with open(self.pkglist, 'r') as pkglist:
             pkgs = pkglist.read()
             pkgs = pkgs.split()
         return pkgs
@@ -403,7 +417,7 @@ class SuseHandler(OsHandler):
         if enterprise:
             self.sources.append('file://' + os.path.join(sourcepath, 'Product-HPC'))
 
-    def prep_root(self):
+    def prep_root(self, args):
         mkdirp(self.targpath)
         if not self.sources:
             targzypp = os.path.join(self.targpath, 'etc/zypp')
@@ -412,6 +426,9 @@ class SuseHandler(OsHandler):
                 '/etc/zypp/repos.d/', os.path.join(targzypp, 'repos.d'))
         for source in self.sources:
             subprocess.check_call(['zypper', '-R', self.targpath, 'ar', source])
+        for source in self.addrepos.split(','):
+            source = 'file://' + source
+            subprocess.check_call(['zypper', '-R', self.targpath, 'ar', source])
         mydir = get_mydir(self.oscategory)
         mkdirp(os.path.join(self.targpath, 'usr/lib/dracut/modules.d'))
         mkdirp(os.path.join(self.targpath, 'etc/dracut.conf.d'))
@@ -437,7 +454,7 @@ class ElHandler(OsHandler):
         self.yumargs.extend(self.list_packages())
 
     def set_source(self, sourcepath):
-        yumconfig = create_yumconf(sourcepath)
+        yumconfig = create_yumconf(sourcepath, self.addrepos)
         self.yumargs.extend(
             ['--setopt=reposdir={0}'.format(yumconfig), '--disablerepo=*',
              '--enablerepo=genimage-*'])
@@ -449,7 +466,7 @@ class ElHandler(OsHandler):
             ['--installroot={0}'.format(targpath),
              '--releasever={0}'.format(self.version), 'install'])
 
-    def prep_root(self):
+    def prep_root(self, args):
         mkdirp(os.path.join(self.targpath, 'usr/lib/dracut/modules.d'))
         mkdirp(os.path.join(self.targpath, 'etc/dracut.conf.d'))
         open(os.path.join(self.targpath, 'etc/resolv.conf'),'w').close()
@@ -557,6 +574,8 @@ def main():
     parser = argparse.ArgumentParser(description='Work with confluent OS cloning and diskless images')
     sps = parser.add_subparsers(dest='subcommand')
     buildp = sps.add_parser('build', help='Build a new diskless image from scratch')
+    buildp.add_argument('-r', '--addrepos', help='Repositories to add in addition to the main source', default='')
+    buildp.add_argument('-p', '--packagelist', help='Filename of package list', default='')
     buildp.add_argument('-s', '--source', help='Directory to pull installation from, typically a subdirectory of /var/lib/confluent/distributions. By default, the repositories for the build system are used.')
     buildp.add_argument('-v', '--volume',
         help='Directory to make available in the build environment.  -v / will '
@@ -642,7 +661,7 @@ def build_root_backend(optargs):
     args, oshandler = optargs
     installroot = args.scratchdir
     _mount_constrained_fs(args, installroot)
-    oshandler.prep_root()
+    oshandler.prep_root(optargs)
 
 
 def _mount_constrained_fs(args, installroot):