From b0b0b612a103d2ffae512857794652260801b303 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Mon, 6 Jan 2020 20:58:46 -0700
Subject: [PATCH] Added Linux mount and unmount functions

* If not running with root priviledges then udevil is used.
---
 scripts/wk.prev/functions/ufd.py | 19 -------------
 scripts/wk/os/linux.py           | 46 +++++++++++++++++++++++++++-----
 2 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/scripts/wk.prev/functions/ufd.py b/scripts/wk.prev/functions/ufd.py
index b5e8cf01..ceae0b75 100644
--- a/scripts/wk.prev/functions/ufd.py
+++ b/scripts/wk.prev/functions/ufd.py
@@ -128,19 +128,6 @@ def is_valid_path(path_obj, path_type):
   return valid_path
 
 
-def mount(mount_source, mount_point, read_write=False):
-  """Mount mount_source on mount_point."""
-  os.makedirs(mount_point, exist_ok=True)
-  cmd = [
-    'mount',
-    mount_source,
-    mount_point,
-    '-o',
-    'rw' if read_write else 'ro',
-    ]
-  run_program(cmd)
-
-
 def prep_device(dev_path, label, use_mbr=False, indent=2):
   """Format device in preparation for applying the WizardKit components
 
@@ -250,12 +237,6 @@ def show_selections(args, sources, ufd_dev, ufd_sources):
   print_standard(' ')
 
 
-def unmount(mount_point):
-  """Unmount mount_point."""
-  cmd = ['umount', mount_point]
-  run_program(cmd)
-
-
 def update_boot_entries(boot_entries, boot_files, iso_label, ufd_label):
   """Update boot files for UFD usage"""
   configs = []
diff --git a/scripts/wk/os/linux.py b/scripts/wk/os/linux.py
index c089ceb8..253b1dd0 100644
--- a/scripts/wk/os/linux.py
+++ b/scripts/wk/os/linux.py
@@ -62,6 +62,27 @@ def make_temp_file():
   return pathlib.Path(proc.stdout.strip())
 
 
+def mount(source, mount_point=None, read_write=False):
+  """Mount source (on mount_point if provided).
+
+  NOTE: If not running_as_root() then udevil will be used.
+  """
+  cmd = [
+    'mount',
+    '-o', 'rw' if read_write else 'ro',
+    source,
+    ]
+  if not running_as_root():
+    cmd.insert(0, 'udevil')
+  if mount_point:
+    cmd.append(mount_point)
+
+  # Run mount command
+  proc = run_program(cmd, check=False)
+  if not proc.returncode == 0:
+    raise RuntimeError(f'Failed to mount: {source} on {mount_point}')
+
+
 def mount_volumes(device_path=None, read_write=False, scan_corestorage=False):
   """Mount all detected volumes, returns list.
 
@@ -117,12 +138,7 @@ def mount_volumes(device_path=None, read_write=False, scan_corestorage=False):
 
     # Attempt to mount volume
     if not already_mounted:
-      cmd = [
-        'udevil',
-        'mount',
-        '-o', 'rw' if read_write else 'ro',
-        vol.path,
-        ]
+      mount(vol.path, read_write=read_write)
       proc = run_program(cmd, check=False)
       if proc.returncode:
         result += 'Failed to mount'
@@ -202,5 +218,23 @@ def scan_corestorage_container(container, timeout=300):
   return inner_volumes
 
 
+def unmount(source_or_mountpoint):
+  """Unmount source_or_mountpoint.
+
+  NOTE: If not running_as_root() then udevil will be used.
+  """
+  cmd = [
+    'umount',
+    source_or_mountpoint,
+    ]
+  if not running_as_root():
+    cmd.insert(0, 'udevil')
+
+  # Run unmount command
+  proc = run_program(cmd, check=False)
+  if not proc.returncode == 0:
+    raise RuntimeError(f'Failed to unmount: {source_or_mountpoint}')
+
+
 if __name__ == '__main__':
   print("This file is not meant to be called directly.")