From bbb4f5a0ecd7d422d1418813a185f9447c033244 Mon Sep 17 00:00:00 2001
From: Jarrod Johnson <jjohnson2@lenovo.com>
Date: Sun, 3 May 2020 19:18:12 -0400
Subject: [PATCH] Add Ubuntu media to osimport

---
 librarian/osimport | 48 +++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 41 insertions(+), 7 deletions(-)

diff --git a/librarian/osimport b/librarian/osimport
index 3c3162c3..31c85617 100644
--- a/librarian/osimport
+++ b/librarian/osimport
@@ -10,6 +10,7 @@ import sys
 COPY = 0
 EXTRACT = 1
 READFILES = set([
+    'README.diskdefines',
     'media.1/products',
     'media.2/products',
     '.discinfo',
@@ -27,7 +28,7 @@ from libarchive.ffi import (
     read_data_block, write_data_block, write_finish_entry, ARCHIVE_EOF
 )
 
-def extract_entries(entries, flags=0, callback=None, totalsize=None):
+def extract_entries(entries, flags=0, callback=None, totalsize=None, extractlist=None):
     """Extracts the given archive entries into the current directory.
     """
     buff, size, offset = c_void_p(), c_size_t(), c_longlong()
@@ -37,6 +38,8 @@ def extract_entries(entries, flags=0, callback=None, totalsize=None):
         for entry in entries:
             if str(entry).endswith('TRANS.TBL'):
                 continue
+            if extractlist and str(entry) not in extractlist:
+                continue
             write_header(write_p, entry._entry_p)
             read_p = entry._archive_p
             while 1:
@@ -49,7 +52,7 @@ def extract_entries(entries, flags=0, callback=None, totalsize=None):
                 write_data_block(write_p, buff, size, offset)
             write_finish_entry(write_p)
 
-def extract_file(filepath, flags=0, callback=lambda x: None, imginfo=()):
+def extract_file(filepath, flags=0, callback=lambda x: None, imginfo=(), extractlist=None):
     """Extracts an archive from a file into the current directory."""
     totalsize = 0
     for img in imginfo:
@@ -57,7 +60,7 @@ def extract_file(filepath, flags=0, callback=lambda x: None, imginfo=()):
             continue
         totalsize += imginfo[img]
     with libarchive.file_reader(filepath) as archive:
-        extract_entries(archive, flags, callback, totalsize)
+        extract_entries(archive, flags, callback, totalsize, extractlist)
 
 def check_centos(isoinfo):
     ver = None
@@ -76,6 +79,36 @@ def check_centos(isoinfo):
         return None
     return {'name': 'centos-{0}-{1}'.format(ver, arch), 'method': EXTRACT}
 
+def check_ubuntu(isoinfo):
+    if 'README.diskdefines' not in isoinfo[1]:
+        return None
+    arch = None
+    variant = None
+    ver = None
+    diskdefs = isoinfo[1]['README.diskdefines']
+    for info in diskdefs.split(b'\n'):
+        if not info:
+            continue
+        _, key, val = info.split(b' ', 2)
+        val = val.strip()
+        if key == b'ARCH':
+            arch = val
+        elif key == b'DISKNAME':
+            variant, ver, _ = val.split(b' ', 2)
+            if variant != b'Ubuntu-Server':
+                return None
+    if variant:
+        if not isinstance(ver, str):
+            ver = ver.decode('utf8')
+        if not isinstance(arch, str):
+            arch = arch.decode('utf8')
+        return {'name': 'ubuntu-{0}-{1}'.format(ver, arch),
+                'method': 'EXTRACTANDCOPY',
+                'extractlist': ['casper/vmlinuz', 'casper/initrd',
+                'EFI/BOOT/BOOTx64.EFI', 'EFI/BOOT/grubx64.efi'
+                ],
+                'copyto': 'install.iso'}
+
 
 def check_sles(isoinfo):
     ver = None
@@ -193,10 +226,11 @@ def import_image(filename):
     os.chdir(targpath)
     print('Importing OS to ' + targpath + ':')
     printit({'progress': 0.0})
-    if identity['method'] == EXTRACT:
-        extract_file(filename, callback=printit, imginfo=imginfo)
-    elif identity['method'] == COPY:
-        targpath = os.path.join(targpath, os.path.basename(filename))
+    if 'EXTRACT' in identity['method']:
+        extract_file(filename, callback=printit, imginfo=imginfo, extractlist=identity.get('extractlist', None))
+    if 'COPY' in identity['method']:
+        basename = identity.get('copyto', os.path.basename(filename))
+        targpath = os.path.join(targpath, basename)
         shutil.copyfile(filename, targpath)
     printit({'progress': 1.0})
     sys.stdout.write('\n')