Moved most logic to functions.ufd
This commit is contained in:
parent
164fd4c646
commit
3310e3d253
GPG key ID:
152FAC923B0E132C
2 changed files with 137 additions and 102 deletions
|
|
@ -7,39 +7,12 @@ import sys
|
||||||
|
|
||||||
# Init
|
# Init
|
||||||
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
|
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
|
||||||
from collections import OrderedDict
|
|
||||||
from docopt import docopt
|
from docopt import docopt
|
||||||
from functions.common import *
|
from functions.common import *
|
||||||
from functions.ufd import *
|
from functions.ufd import *
|
||||||
from settings.ufd import *
|
from settings.ufd import *
|
||||||
init_global_vars(silent=True)
|
init_global_vars(silent=True)
|
||||||
|
|
||||||
# Functions
|
|
||||||
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
|
|
||||||
|
|
||||||
|
|
||||||
# Main section
|
# Main section
|
||||||
if __name__ == '__main__':
|
if __name__ == '__main__':
|
||||||
# Set log
|
# Set log
|
||||||
|
|
@ -56,7 +29,7 @@ if __name__ == '__main__':
|
||||||
print_standard(' ')
|
print_standard(' ')
|
||||||
|
|
||||||
# Check if running as root
|
# Check if running as root
|
||||||
if os.geteuid() != 0:
|
if not running_as_root():
|
||||||
print_error('ERROR: This script is meant to be run as root.')
|
print_error('ERROR: This script is meant to be run as root.')
|
||||||
abort(False)
|
abort(False)
|
||||||
|
|
||||||
|
|
@ -71,81 +44,12 @@ if __name__ == '__main__':
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Verify selections
|
# Verify selections
|
||||||
sources = OrderedDict()
|
ufd_dev = verify_ufd(args['--ufd-device'])
|
||||||
## UFD
|
sources = verify_sources(args, UFD_SOURCES)
|
||||||
try:
|
show_selections(args, sources, ufd_dev, UFD_SOURCES)
|
||||||
ufd_dev = find_path(args['--ufd-device'])
|
confirm_selections(args)
|
||||||
except FileNotFoundError:
|
|
||||||
print_error('ERROR: UFD device not found: {}'.format(
|
|
||||||
args['--ufd-device']))
|
|
||||||
abort(False)
|
|
||||||
if not is_valid_path(ufd_dev, 'UFD'):
|
|
||||||
print_error('ERROR: Invalid UFD device: {}'.format(ufd_dev))
|
|
||||||
abort(False)
|
|
||||||
## Sources
|
|
||||||
for label, data in UFD_SOURCES.items():
|
|
||||||
s_path = args[data['Arg']]
|
|
||||||
if s_path:
|
|
||||||
try:
|
|
||||||
s_path_obj = find_path(s_path)
|
|
||||||
except FileNotFoundError:
|
|
||||||
print_error('ERROR: {} not found: {}'.format(label, s_path))
|
|
||||||
abort(False)
|
|
||||||
if not is_valid_path(s_path_obj, data['Type']):
|
|
||||||
print_error('ERROR: Invalid {} source: {}'.format(label, s_path))
|
|
||||||
abort(False)
|
|
||||||
sources[label] = s_path_obj
|
|
||||||
|
|
||||||
# Show selections
|
# TODO: DELETEME
|
||||||
## Sources
|
|
||||||
print_info('Sources')
|
|
||||||
for label in UFD_SOURCES.keys():
|
|
||||||
if label in sources:
|
|
||||||
print_standard(' {label:<18} {path}'.format(
|
|
||||||
label=label+':',
|
|
||||||
path=sources[label],
|
|
||||||
))
|
|
||||||
else:
|
|
||||||
print_standard(' {label:<18} {YELLOW}Not Specified{CLEAR}'.format(
|
|
||||||
label=label+':',
|
|
||||||
**COLORS,
|
|
||||||
))
|
|
||||||
## Destination
|
|
||||||
print_standard(' ')
|
|
||||||
print_info('Destination')
|
|
||||||
cmd = [
|
|
||||||
'lsblk', '--nodeps', '--noheadings',
|
|
||||||
'--output', 'NAME,TRAN,SIZE,VENDOR,MODEL,SERIAL',
|
|
||||||
]
|
|
||||||
result = run_program(cmd, check=False, encoding='utf-8', errors='ignore')
|
|
||||||
print_standard(result.stdout.strip())
|
|
||||||
cmd = [
|
|
||||||
'lsblk', '--noheadings',
|
|
||||||
'--output', 'NAME,SIZE,FSTYPE,LABEL,MOUNTPOINT',
|
|
||||||
]
|
|
||||||
result = run_program(cmd, check=False, encoding='utf-8', errors='ignore')
|
|
||||||
for line in result.stdout.splitlines()[1:]:
|
|
||||||
print_standard(line)
|
|
||||||
if args['--update']:
|
|
||||||
print_warning('Updating kit in-place')
|
|
||||||
elif args['--use-mbr']:
|
|
||||||
print_warning('Formatting using legacy MBR')
|
|
||||||
print_standard(' ')
|
|
||||||
if not ask('Is the above information correct?'):
|
|
||||||
abort(False)
|
|
||||||
## Safety check
|
|
||||||
if not args['--update']:
|
|
||||||
print_standard(' ')
|
|
||||||
print_warning('SAFETY CHECK')
|
|
||||||
print_standard(
|
|
||||||
'All data will be DELETED from the disk and partition(s) listed above.')
|
|
||||||
print_standard(
|
|
||||||
'This is irreversible and will lead to {RED}DATA LOSS.{CLEAR}'.format(
|
|
||||||
**COLORS))
|
|
||||||
if not ask('Asking again to confirm, is this correct?'):
|
|
||||||
abort(False)
|
|
||||||
|
|
||||||
print_standard(' ')
|
|
||||||
print_success("It's go-time!")
|
print_success("It's go-time!")
|
||||||
exit_script()
|
exit_script()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -4,6 +4,8 @@ import os
|
||||||
import re
|
import re
|
||||||
import shutil
|
import shutil
|
||||||
import pathlib
|
import pathlib
|
||||||
|
from collections import OrderedDict
|
||||||
|
from functions.common import *
|
||||||
|
|
||||||
|
|
||||||
def case_insensitive_search(path, item):
|
def case_insensitive_search(path, item):
|
||||||
|
|
@ -35,6 +37,25 @@ def case_insensitive_search(path, item):
|
||||||
return real_path
|
return real_path
|
||||||
|
|
||||||
|
|
||||||
|
def confirm_selections(args):
|
||||||
|
"""Ask tech to confirm selections, twice if necessary."""
|
||||||
|
if not ask('Is the above information correct?'):
|
||||||
|
abort(False)
|
||||||
|
## Safety check
|
||||||
|
if not args['--update']:
|
||||||
|
print_standard(' ')
|
||||||
|
print_warning('SAFETY CHECK')
|
||||||
|
print_standard(
|
||||||
|
'All data will be DELETED from the disk and partition(s) listed above.')
|
||||||
|
print_standard(
|
||||||
|
'This is irreversible and will lead to {RED}DATA LOSS.{CLEAR}'.format(
|
||||||
|
**COLORS))
|
||||||
|
if not ask('Asking again to confirm, is this correct?'):
|
||||||
|
abort(False)
|
||||||
|
|
||||||
|
print_standard(' ')
|
||||||
|
|
||||||
|
|
||||||
def find_path(path):
|
def find_path(path):
|
||||||
"""Find path case-insensitively, returns pathlib.Path obj."""
|
"""Find path case-insensitively, returns pathlib.Path obj."""
|
||||||
path_obj = pathlib.Path(path).resolve()
|
path_obj = pathlib.Path(path).resolve()
|
||||||
|
|
@ -62,6 +83,31 @@ def find_path(path):
|
||||||
return path_obj
|
return path_obj
|
||||||
|
|
||||||
|
|
||||||
|
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 is_valid_path(path_obj, path_type):
|
def is_valid_path(path_obj, path_type):
|
||||||
"""Verify path_obj is valid by type, returns bool."""
|
"""Verify path_obj is valid by type, returns bool."""
|
||||||
valid_path = False
|
valid_path = False
|
||||||
|
|
@ -126,6 +172,91 @@ def recursive_copy(source, dest, overwrite=True):
|
||||||
raise FileExistsError('Refusing to delete file: {}'.format(dest))
|
raise FileExistsError('Refusing to delete file: {}'.format(dest))
|
||||||
|
|
||||||
|
|
||||||
|
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."""
|
||||||
|
|
||||||
|
# Sources
|
||||||
|
print_info('Sources')
|
||||||
|
for label in ufd_sources.keys():
|
||||||
|
if label in sources:
|
||||||
|
print_standard(' {label:<18} {path}'.format(
|
||||||
|
label=label+':',
|
||||||
|
path=sources[label],
|
||||||
|
))
|
||||||
|
else:
|
||||||
|
print_standard(' {label:<18} {YELLOW}Not Specified{CLEAR}'.format(
|
||||||
|
label=label+':',
|
||||||
|
**COLORS,
|
||||||
|
))
|
||||||
|
print_standard(' ')
|
||||||
|
|
||||||
|
# Destination
|
||||||
|
print_info('Destination')
|
||||||
|
cmd = [
|
||||||
|
'lsblk', '--nodeps', '--noheadings',
|
||||||
|
'--output', 'NAME,TRAN,SIZE,VENDOR,MODEL,SERIAL',
|
||||||
|
]
|
||||||
|
result = run_program(cmd, check=False, encoding='utf-8', errors='ignore')
|
||||||
|
print_standard(result.stdout.strip())
|
||||||
|
cmd = [
|
||||||
|
'lsblk', '--noheadings',
|
||||||
|
'--output', 'NAME,SIZE,FSTYPE,LABEL,MOUNTPOINT',
|
||||||
|
]
|
||||||
|
result = run_program(cmd, check=False, encoding='utf-8', errors='ignore')
|
||||||
|
for line in result.stdout.splitlines()[1:]:
|
||||||
|
print_standard(line)
|
||||||
|
|
||||||
|
# Notes
|
||||||
|
if args['--update']:
|
||||||
|
print_warning('Updating kit in-place')
|
||||||
|
elif args['--use-mbr']:
|
||||||
|
print_warning('Formatting using legacy MBR')
|
||||||
|
print_standard(' ')
|
||||||
|
|
||||||
|
|
||||||
|
def verify_sources(args, ufd_sources):
|
||||||
|
"""Check all sources and abort if necessary, returns dict."""
|
||||||
|
sources = OrderedDict()
|
||||||
|
|
||||||
|
for label, data in ufd_sources.items():
|
||||||
|
s_path = args[data['Arg']]
|
||||||
|
if s_path:
|
||||||
|
try:
|
||||||
|
s_path_obj = find_path(s_path)
|
||||||
|
except FileNotFoundError:
|
||||||
|
print_error('ERROR: {} not found: {}'.format(label, s_path))
|
||||||
|
abort(False)
|
||||||
|
if not is_valid_path(s_path_obj, data['Type']):
|
||||||
|
print_error('ERROR: Invalid {} source: {}'.format(label, s_path))
|
||||||
|
abort(False)
|
||||||
|
sources[label] = s_path_obj
|
||||||
|
|
||||||
|
return sources
|
||||||
|
|
||||||
|
|
||||||
|
def verify_ufd(dev_path):
|
||||||
|
"""Check that dev_path is a valid UFD, returns pathlib.Path obj."""
|
||||||
|
ufd_dev = None
|
||||||
|
|
||||||
|
try:
|
||||||
|
ufd_dev = find_path(dev_path)
|
||||||
|
except FileNotFoundError:
|
||||||
|
print_error('ERROR: UFD device not found: {}'.format(
|
||||||
|
args['--ufd-device']))
|
||||||
|
abort(False)
|
||||||
|
|
||||||
|
if not is_valid_path(ufd_dev, 'UFD'):
|
||||||
|
print_error('ERROR: Invalid UFD device: {}'.format(ufd_dev))
|
||||||
|
abort(False)
|
||||||
|
|
||||||
|
return ufd_dev
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
if __name__ == '__main__':
|
||||||
print("This file is not meant to be called directly.")
|
print("This file is not meant to be called directly.")
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue