Merge pull request #267 from postmarketOS/fix/242-validate-input

Fix #242: Validate input in pmb.helpers.cli

pmb/chroot/initfs.py

@@ -48,8 +48,8 @@ def extract(args, flavor, suffix, log_message=False):
     inside = "/tmp/initfs-extracted"
     outside = args.work + "/chroot_" + suffix + inside
     if os.path.exists(outside):
-        if pmb.helpers.cli.ask(args, "Extraction folder " + outside +
-                               " already exists. Do you want to overwrite it?") != "y":
+        if not pmb.helpers.cli.confirm(args, "Extraction folder " + outside +
+                                       " already exists. Do you want to overwrite it?"):
             raise RuntimeError("Aborted!")
         pmb.chroot.root(args, ["rm", "-r", inside], suffix)
 

pmb/chroot/zap.py

@@ -42,5 +42,5 @@ def zap(args):
         pattern = os.path.abspath(args.work + "/" + pattern)
         matches = glob.glob(pattern)
         for match in matches:
-            if pmb.helpers.cli.ask(args, "Remove " + match + "?") == "y":
+            if pmb.helpers.cli.confirm(args, "Remove " + match + "?"):
                 pmb.helpers.run.root(args, ["rm", "-rf", match])

pmb/config/init.py

@@ -26,6 +26,38 @@ import pmb.helpers.devices
 import pmb.helpers.ui
 
 
+def ask_for_work_path(args):
+    """
+    Ask for the work path, until we can create it (when it does not exist) and
+    write into it.
+    :returns: the work path
+    """
+    logging.info("Location of the 'work' path. Multiple chroots"
+                 " (native, device arch, device rootfs) will be created"
+                 " in there.")
+    while True:
+        try:
+            ret = os.path.expanduser(pmb.helpers.cli.ask(
+                args, "Work path", None, args.work, False))
+            os.makedirs(ret + "/chroot_native", 0o700, True)
+            return ret
+        except OSError:
+            logging.fatal("ERROR: Could not create this folder, or write"
+                          " inside it! Please try again.")
+
+
+def ask_for_ui(args):
+    ui_list = pmb.helpers.ui.list(args)
+    logging.info("Available user interfaces (" +
+                 str(len(ui_list) - 1) + "): " + ", ".join(ui_list))
+    while True:
+        ret = pmb.helpers.cli.ask(args, "User interface", None, args.ui, True)
+        if ret in ui_list:
+            return ret
+        logging.fatal("ERROR: Invalid user interface specified, please type in"
+                      " one from the list above.")
+
+
 def init(args):
     cfg = pmb.config.load(args)
 
@@ -36,20 +68,11 @@ def init(args):
     logging.info("Available (" + str(len(devices)) + "): " +
                  ", ".join(devices))
     cfg["pmbootstrap"]["device"] = pmb.helpers.cli.ask(args, "Device",
-                                                       None, args.device)
+                                                       None, args.device, False, "[a-z0-9]+-[a-z0-9]+")
 
-    # UI selection
-    ui_list = pmb.helpers.ui.list(args)
-    logging.info("Available user interfaces (" + str(len(ui_list) - 1) + "): " + ", ".join(ui_list))
-    cfg["pmbootstrap"]["ui"] = pmb.helpers.cli.ask(args, "User Interface:",
-                                                   None, args.ui, True)
-
-    # Work folder
-    logging.info("Location of the 'work' path. Multiple chroots (native,"
-                 " device arch, device rootfs) will be created in there.")
-    cfg["pmbootstrap"]["work"] = os.path.expanduser(pmb.helpers.cli.ask(args, "Work path",
-                                                                        None, args.work, False))
-    os.makedirs(cfg["pmbootstrap"]["work"], 0o700, True)
+    # UI and work folder
+    cfg["pmbootstrap"]["ui"] = ask_for_ui(args)
+    cfg["pmbootstrap"]["work"] = ask_for_work_path(args)
 
     # Parallel job count
     default = args.jobs
@@ -58,17 +81,15 @@ def init(args):
     logging.info("How many jobs should run parallel on this machine, when"
                  " compiling?")
     cfg["pmbootstrap"]["jobs"] = pmb.helpers.cli.ask(args, "Jobs",
-                                                     None, default)
+                                                     None, default, validation_regex="[1-9][0-9]*")
+
     # Timestamp based rebuilds
-    default = "y"
-    if not args.timestamp_based_rebuild:
-        default = "n"
     logging.info("Rebuild packages, when the last modified timestamp changed,"
                  " even if the version did not change? This makes pmbootstrap"
                  " behave more like 'make'.")
-    answer = pmb.helpers.cli.ask(args, "Timestamp based rebuilds",
-                                 default=default)
-    cfg["pmbootstrap"]["timestamp_based_rebuild"] = str(answer == "y")
+    answer = pmb.helpers.cli.confirm(args, "Timestamp based rebuilds",
+                                     default=args.timestamp_based_rebuild)
+    cfg["pmbootstrap"]["timestamp_based_rebuild"] = str(answer)
 
     # Do not save aports location to config file
     del cfg["pmbootstrap"]["aports"]

pmb/helpers/cli.py

@@ -17,23 +17,51 @@ You should have received a copy of the GNU General Public License
 along with pmbootstrap.  If not, see <http://www.gnu.org/licenses/>.
 """
 import datetime
+import logging
+import re
 
 
-def ask(args, question="Continue?", choices=['y', 'n'], default='n',
-        lowercase_answer=True):
-    date = datetime.datetime.now().strftime("%H:%M:%S")
-    question = "[" + date + "] " + question
-    if choices:
-        question += " (" + str.join("/", choices) + ")"
-    if default:
-        question += " [" + str(default) + "]"
+def ask(args, question="Continue?", choices=["y", "n"], default="n",
+        lowercase_answer=True, validation_regex=None):
+    """
+    Ask a question on the terminal. When validation_regex is set, the user gets
+    asked until the answer matches the regex.
+    :returns: the user's answer
+    """
+    while True:
+        date = datetime.datetime.now().strftime("%H:%M:%S")
+        question_full = "[" + date + "] " + question
+        if choices:
+            question_full += " (" + str.join("/", choices) + ")"
+        if default:
+            question_full += " [" + str(default) + "]"
 
-    ret = input(question + ": ")
-    if lowercase_answer:
-        ret = ret.lower()
-    if ret == "":
-        ret = str(default)
+        ret = input(question_full + ": ")
+        if lowercase_answer:
+            ret = ret.lower()
+        if ret == "":
+            ret = str(default)
 
-    args.logfd.write(question + " " + ret + "\n")
-    args.logfd.flush()
-    return ret
+        args.logfd.write(question_full + " " + ret + "\n")
+        args.logfd.flush()
+
+        # Validate with regex
+        if not validation_regex:
+            return ret
+
+        pattern = re.compile(validation_regex)
+        if pattern.match(ret):
+            return ret
+
+        logging.fatal("ERROR: Input did not pass validation (regex: " +
+                      validation_regex + "). Please try again.")
+
+
+def confirm(args, question="Continue?", default=False):
+    """
+    Convenience wrapper around ask for simple yes-no questions with validation.
+    :returns: True for "y", False for "n"
+    """
+    default_str = "y" if default else "n"
+    answer = ask(args, question, ["y", "n"], default_str, True, "(y|n)")
+    return answer == "y"

pmb/install/blockdevice.py

@@ -37,8 +37,8 @@ def mount_sdcard(args):
         if pmb.helpers.mount.ismount(path):
             raise RuntimeError(path + " is mounted! We will not attempt"
                                " to format this!")
-    if pmb.helpers.cli.ask(args, "EVERYTHING ON " + args.sdcard + " WILL BE"
-                           " ERASED! CONTINUE?") != "y":
+    if not pmb.helpers.cli.confirm(args, "EVERYTHING ON " + args.sdcard +
+                                   " WILL BE ERASED! CONTINUE?"):
         raise RuntimeError("Aborted.")
 
     logging.info("(native) mount /dev/install (host: " + args.sdcard + ")")
@@ -66,7 +66,7 @@ def create_and_mount_image(args):
     logging.info("(native) create " + args.device + ".img (" + size + ")")
     logging.info("WARNING: Make sure, that your target device's partition"
                  " table has allocated at least " + size + " as system partition!")
-    if pmb.helpers.cli.ask(args) != "y":
+    if not pmb.helpers.cli.confirm(args):
         raise RuntimeError("Aborted.")
 
     pmb.chroot.user(args, ["mkdir", "-p", "/home/user/rootfs"])