Updated common.py

This commit is contained in:
2Shirt 2018-11-23 18:06:57 -07:00
parent 923ddd7cd3
commit 18b13cf506
Signed by: 2Shirt
GPG key ID: 152FAC923B0E132C

View file

@ -32,7 +32,8 @@ COLORS = {
'BLUE': '\033[34m'
}
try:
HKU = winreg.HKEY_USERS
HKU = winreg.HKEY_USERS
HKCR = winreg.HKEY_CLASSES_ROOT
HKCU = winreg.HKEY_CURRENT_USER
HKLM = winreg.HKEY_LOCAL_MACHINE
except NameError:
@ -167,14 +168,13 @@ def exit_script(return_value=0):
# Remove dirs (if empty)
for dir in ['BackupDir', 'LogDir', 'TmpDir']:
try:
dir = global_vars[dir]
os.rmdir(dir)
os.rmdir(global_vars[dir])
except Exception:
pass
# Open Log (if it exists)
log = global_vars.get('LogFile', '')
if log and os.path.exists(log) and psutil.WINDOWS:
if log and os.path.exists(log) and psutil.WINDOWS and ENABLED_OPEN_LOGS:
try:
extract_item('NotepadPlusPlus', silent=True)
popen_program(
@ -499,6 +499,8 @@ def sleep(seconds=2):
def stay_awake():
"""Prevent the system from sleeping or hibernating."""
# DISABLED due to VCR2008 dependency
return
# Bail if caffeine is already running
for proc in psutil.process_iter():
if proc.name() == 'caffeine.exe':
@ -506,7 +508,7 @@ def stay_awake():
# Extract and run
extract_item('Caffeine', silent=True)
try:
popen_program(global_vars['Tools']['Caffeine'])
popen_program([global_vars['Tools']['Caffeine']])
except Exception:
print_error('ERROR: No caffeine available.')
print_warning('Please set the power setting to High Performance.')
@ -601,9 +603,10 @@ global_vars: {}'''.format(f.read(), sys.argv, global_vars)
CRASH_SERVER['Url'],
global_vars.get('Date-Time', 'Unknown Date-Time'),
filename)
r = requests.put(url, data=data,
headers = {'X-Requested-With': 'XMLHttpRequest'},
auth = (CRASH_SERVER['User'], CRASH_SERVER['Pass']))
r = requests.put(
url, data=data,
headers={'X-Requested-With': 'XMLHttpRequest'},
auth=(CRASH_SERVER['User'], CRASH_SERVER['Pass']))
# Raise exception if upload NS
if not r.ok:
raise Exception
@ -752,6 +755,9 @@ def make_tmp_dirs():
"""Make temp directories."""
os.makedirs(global_vars['BackupDir'], exist_ok=True)
os.makedirs(global_vars['LogDir'], exist_ok=True)
os.makedirs(r'{}\{}'.format(
global_vars['LogDir'], KIT_NAME_FULL), exist_ok=True)
os.makedirs(r'{}\Tools'.format(global_vars['LogDir']), exist_ok=True)
os.makedirs(global_vars['TmpDir'], exist_ok=True)
def set_common_vars():
@ -767,11 +773,9 @@ def set_common_vars():
**global_vars)
global_vars['ClientDir'] = r'{SYSTEMDRIVE}\{prefix}'.format(
prefix=KIT_NAME_SHORT, **global_vars['Env'])
global_vars['BackupDir'] = r'{ClientDir}\Backups\{Date}'.format(
global_vars['BackupDir'] = r'{ClientDir}\Backups'.format(
**global_vars)
global_vars['LogDir'] = r'{ClientDir}\Info\{Date}'.format(
**global_vars)
global_vars['ProgBackupDir'] = r'{ClientDir}\Backups'.format(
global_vars['LogDir'] = r'{ClientDir}\Logs\{Date}'.format(
**global_vars)
global_vars['QuarantineDir'] = r'{ClientDir}\Quarantine'.format(
**global_vars)
@ -794,5 +798,12 @@ def set_linux_vars():
'SevenZip': '7z',
}
def set_log_file(log_name):
"""Sets global var LogFile and creates path as needed."""
folder_path = r'{}\{}'.format(global_vars['LogDir'], KIT_NAME_FULL)
log_file = r'{}\{}'.format(folder_path, log_name)
os.makedirs(folder_path, exist_ok=True)
global_vars['LogFile'] = log_file
if __name__ == '__main__':
print("This file is not meant to be called directly.")