From c135d686df207ce09b9dc8a227ece96d584c5184 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Mon, 6 Jan 2020 20:27:59 -0700
Subject: [PATCH] Added Linux functions for building UFDs

---
 scripts/wk.prev/functions/ufd.py | 30 ----------------------
 scripts/wk/os/linux.py           | 44 ++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 30 deletions(-)

diff --git a/scripts/wk.prev/functions/ufd.py b/scripts/wk.prev/functions/ufd.py
index ae2e27ca..b5e8cf01 100644
--- a/scripts/wk.prev/functions/ufd.py
+++ b/scripts/wk.prev/functions/ufd.py
@@ -71,31 +71,6 @@ def find_first_partition(dev_path):
   return part_path
 
 
-def get_user_home(user):
-  """Get path to user's home dir, returns str."""
-  home_dir = None
-  cmd = ['getent', 'passwd', user]
-  result = run_program(cmd, encoding='utf-8', errors='ignore', check=False)
-  try:
-    home_dir = result.stdout.split(':')[5]
-  except Exception:
-    # Just use HOME from ENV (or '/root' if that fails)
-    home_dir = os.environ.get('HOME', '/root')
-
-  return home_dir
-
-
-def get_user_name():
-  """Get real user name, returns str."""
-  user = None
-  if 'SUDO_USER' in os.environ:
-    user = os.environ.get('SUDO_USER', 'Unknown')
-  else:
-    user = os.environ.get('USER', 'Unknown')
-
-  return user
-
-
 def hide_items(ufd_dev, items):
   """Set FAT32 hidden flag for items."""
   # pylint: disable=invalid-name
@@ -231,11 +206,6 @@ def remove_arch():
   shutil.rmtree(find_path('/mnt/UFD/arch'))
 
 
-def running_as_root():
-  """Check if running with effective UID of 0, returns bool."""
-  return os.geteuid() == 0
-
-
 def show_selections(args, sources, ufd_dev, ufd_sources):
   """Show selections including non-specified options."""
 
diff --git a/scripts/wk/os/linux.py b/scripts/wk/os/linux.py
index 14e6015b..c089ceb8 100644
--- a/scripts/wk/os/linux.py
+++ b/scripts/wk/os/linux.py
@@ -2,6 +2,7 @@
 # vim: sts=2 sw=2 ts=2
 
 import logging
+import os
 import pathlib
 import re
 import subprocess
@@ -17,6 +18,44 @@ UUID_CORESTORAGE = '53746f72-6167-11aa-aa11-00306543ecac'
 
 
 # Functions
+def get_user_home(user):
+  """Get path to user's home dir, returns pathlib.Path obj."""
+  home = None
+
+  # Get path from user details
+  cmd = ['getent', 'passwd', user]
+  proc = run_program(cmd, check=False)
+  try:
+    home = proc.stdout.split(':')[5]
+  except IndexError:
+    # Try using environment variable
+    home = os.environ.get('HOME')
+
+  # Raise exception if necessary
+  if not home:
+    raise RuntimeError(f'Failed to find home for: {user}')
+
+  # Done
+  return pathlib.Path(home)
+
+
+def get_user_name():
+  """Get real user name, returns str."""
+  user = None
+
+  # Query environment
+  user = os.environ.get('SUDO_USER')
+  if not user:
+    user = os.environ.get('USER')
+
+  # Raise exception if necessary
+  if not user:
+    raise RuntimeError('Failed to determine user')
+
+  # Done
+  return user
+
+
 def make_temp_file():
   """Make temporary file, returns pathlib.Path() obj."""
   proc = run_program(['mktemp'], check=False)
@@ -111,6 +150,11 @@ def mount_volumes(device_path=None, read_write=False, scan_corestorage=False):
   return report
 
 
+def running_as_root():
+  """Check if running with effective UID of 0, returns bool."""
+  return os.geteuid() == 0
+
+
 def scan_corestorage_container(container, timeout=300):
   """Scan CoreStorage container for inner volumes, returns list."""
   # TODO: Test Scanning CoreStorage containers