From 4c33c110b758d7fbb9b9ef0c58e8d1a7d6a37d7e Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Mon, 8 Apr 2019 00:35:53 -0700
Subject: [PATCH] Fixed source validation

---
 .bin/Scripts/build-ufd        | 39 +++++++++++++++++++++--------------
 .bin/Scripts/functions/ufd.py | 16 +++++++++++---
 .bin/Scripts/settings/ufd.py  | 28 ++++++++++++-------------
 3 files changed, 50 insertions(+), 33 deletions(-)

diff --git a/.bin/Scripts/build-ufd b/.bin/Scripts/build-ufd
index a5577907..9c048dea 100755
--- a/.bin/Scripts/build-ufd
+++ b/.bin/Scripts/build-ufd
@@ -7,6 +7,7 @@ import sys
 
 # Init
 sys.path.append(os.path.dirname(os.path.realpath(__file__)))
+from collections import OrderedDict
 from docopt import docopt
 from functions.ufd import *
 from settings.ufd import *
@@ -18,33 +19,39 @@ set_log_file('Build UFD ({Date-Time}).log'.format(**global_vars))
 if __name__ == '__main__':
   try:
     args = docopt(DOCSTRING)
-    sources = ()
+    sources = OrderedDict()
 
     # Verify selections
     ## UFD
     try:
-      ufd_dev = get_full_path(args['--ufd-device'])
-      if not is_block_device(ufd_dev):
-        print_error('Invalid UFD device: {}'.format(ufd_dev))
-        abort()
-    except Exception:
-      # TODO Catch FileNotFound exception and abort accordingly
-      raise
+      ufd_dev = find_path(args['--ufd-device'])
+    except FileNotFoundError:
+      print_error('ERROR: UFD device not found: {}'.format(
+        args['--ufd-device']))
+      abort()
+    if not is_valid_path(ufd_dev, 'UFD'):
+      print_error('ERROR: Invalid UFD device: {}'.format(ufd_dev))
+      abort()
     ## Sources
-    for label, source in UFD_SOURCES:
-      if args[source]:
+    for label, data in UFD_SOURCES.items():
+      s_path = args[data['Arg']]
+      if s_path:
         try:
-          sources.append((label, get_full_path(source)))
-        except Exception:
-          # TODO Catch FileNotFound exception and abort accordingly
-          raise
+          s_path_obj = find_path(s_path)
+        except FileNotFoundError:
+          print_error('ERROR: {} not found: {}'.format(label, s_path))
+          abort()
+        if not is_valid_path(s_path_obj, data['Type']):
+          print_error('ERROR: Invalid {} source: {}'.format(label, s_path))
+          abort()
+        sources[label] = s_path_obj
 
     # Show selections
     # TODO FIXME
     print_standard('UFD: {}'.format(ufd_dev))
     print_standard('Sources:')
-    for s in sources:
-      print_standard('  {}: {}'.format(*s))
+    for label, s_path in sources.items():
+      print_standard('  {:<16} {}'.format(label+':', s_path))
 
     # Double-check if formating device
 
diff --git a/.bin/Scripts/functions/ufd.py b/.bin/Scripts/functions/ufd.py
index f0ce6ad3..ae1bc37d 100644
--- a/.bin/Scripts/functions/ufd.py
+++ b/.bin/Scripts/functions/ufd.py
@@ -60,9 +60,19 @@ def find_path(path):
   return path_obj
 
 
-def is_valid_main_kit(path_obj):
-  """Verify PathObj contains the main kit, returns bool."""
-  return path_obj.is_dir() and path_obj.joinpath('.bin').exists()
+def is_valid_path(path_obj, path_type):
+  """Verify path_obj is valid by type, returns bool."""
+  valid_path = False
+  if path_type == 'DIR':
+    valid_path = path_obj.is_dir()
+  elif path_type == 'KIT':
+    valid_path = path_obj.is_dir() and path_obj.joinpath('.bin').exists()
+  elif path_type == 'ISO':
+    valid_path = path_obj.is_file() and path_obj.suffix.lower() == '.iso'
+  elif path_type == 'UFD':
+    valid_path = path_obj.is_block_device()
+
+  return valid_path
 
 
 if __name__ == '__main__':
diff --git a/.bin/Scripts/settings/ufd.py b/.bin/Scripts/settings/ufd.py
index b7f552eb..50fc6934 100644
--- a/.bin/Scripts/settings/ufd.py
+++ b/.bin/Scripts/settings/ufd.py
@@ -1,25 +1,26 @@
 # Wizard Kit: Settings - UFD
 
+from collections import OrderedDict
 from settings.main import *
 
 # General
 DOCSTRING = '''WizardKit: Build UFD
 
 Usage:
-  build-ufd [options] --ufd-device PATH --linux-iso PATH
-            [--linux-minimal-iso PATH]
+  build-ufd [options] --ufd-device PATH --linux PATH
+            [--linux-minimal PATH]
             [--main-kit PATH]
-            [--winpe-iso PATH]
+            [--winpe PATH]
             [--extra-dir PATH]
   build-ufd (-h | --help)
 
 Options:
   -e PATH, --extra-dir PATH
   -k PATH, --main-kit PATH
-  -l PATH, --linux-iso PATH
-  -m PATH, --linux-minimal-iso PATH
+  -l PATH, --linux PATH
+  -m PATH, --linux-minimal PATH
   -u PATH, --ufd-device PATH
-  -w PATH, --winpe-iso PATH
+  -w PATH, --winpe PATH
 
   -d --debug            Enable debug mode
   -h --help             Show this page
@@ -30,14 +31,13 @@ Options:
 '''
 ISO_LABEL = '{}_LINUX'.format(KIT_NAME_SHORT)
 UFD_LABEL = '{}_UFD'.format(KIT_NAME_SHORT)
-UFD_SOURCES = (
-  # NOTE: Using tuple of tuples to ensure copy order
-  ('Linux', '--linux-iso'),
-  ('Linux (Minimal)', '--linux-minimal-iso'),
-  ('WinPE', '--winpe-iso'),
-  ('Main Kit', '--main-kit'),
-  ('Extras', '--extra-dir'),
-  )
+UFD_SOURCES = ({
+  'Linux':            {'Arg': '--linux',          'Type': 'ISO'},
+  'Linux (Minimal)':  {'Arg': '--linux-minimal',  'Type': 'ISO'},
+  'WinPE':            {'Arg': '--winpe',          'Type': 'ISO'},
+  'Main Kit':         {'Arg': '--main-kit',       'Type': 'KIT'},
+  'Extra Dir':        {'Arg': '--extra-dir',      'Type': 'DIR'},
+  })
 
 # Definitions: Boot entries
 ## NOTE: if key path exists uncomment #value# lines