Merge branch 'python-overhaul' into dev
This commit is contained in:
Alan Mason
commit
1b84b1dd91
18 changed files with 2599 additions and 1335 deletions
|
|
@ -518,18 +518,18 @@ if ($MyInvocation.InvocationName -ne ".") {
|
|||
)
|
||||
Start-Process -FilePath $DISM -ArgumentList $ArgumentList -NoNewWindow -Wait
|
||||
|
||||
# Add WK tools
|
||||
# Add tools
|
||||
Write-Host "Copying tools..."
|
||||
Copy-Item -Path "$Root\WK\$Arch" -Destination "$Mount\WK" -Recurse -Force
|
||||
Copy-Item -Path "$Root\WK\_include\*" -Destination "$Mount\WK" -Recurse -Force
|
||||
Copy-Item -Path "$Root\WK\$Arch" -Destination "$Mount\.bin" -Recurse -Force
|
||||
Copy-Item -Path "$Root\WK\_include\*" -Destination "$Mount\.bin" -Recurse -Force
|
||||
if ($Arch -eq "amd64") {
|
||||
$DestIni = "$Mount\WK\HWiNFO\HWiNFO64.INI"
|
||||
$DestIni = "$Mount\.bin\HWiNFO\HWiNFO64.INI"
|
||||
} else {
|
||||
$DestIni = "$Mount\WK\HWiNFO\HWiNFO32.INI"
|
||||
$DestIni = "$Mount\.bin\HWiNFO\HWiNFO32.INI"
|
||||
}
|
||||
Move-Item -Path "$Mount\WK\HWiNFO\HWiNFO.INI" -Destination $DestIni -Force
|
||||
Copy-Item -Path "$Root\WinPE.jpg" -Destination "$Mount\WK\ConEmu\ConEmu.jpg" -Recurse -Force
|
||||
Copy-Item -Path "$Root\Scripts" -Destination "$Mount\WK\Scripts" -Recurse -Force
|
||||
Move-Item -Path "$Mount\.bin\HWiNFO\HWiNFO.INI" -Destination $DestIni -Force
|
||||
Copy-Item -Path "$Root\WinPE.jpg" -Destination "$Mount\.bin\ConEmu\ConEmu.jpg" -Recurse -Force
|
||||
Copy-Item -Path "$Root\Scripts" -Destination "$Mount\.bin\Scripts" -Recurse -Force
|
||||
|
||||
# Add System32 items
|
||||
$HostSystem32 = "{0}\System32" -f $Env:SystemRoot
|
||||
|
|
@ -557,17 +557,16 @@ if ($MyInvocation.InvocationName -ne ".") {
|
|||
$RegKey = $Hive.OpenSubKey($RegPath)
|
||||
$CurValue = $RegKey.GetValue(
|
||||
"Path", $false, [Microsoft.Win32.RegistryValueOptions]::DoNotExpandEnvironmentNames)
|
||||
$NewValue = "$CurValue;%SystemDrive%\WK\7-Zip;%SystemDrive%\WK\python;%SystemDrive%\WK\wimlib"
|
||||
$NewValue = "$CurValue;%SystemDrive%\.bin\7-Zip;%SystemDrive%\.bin\python;%SystemDrive%\.bin\wimlib"
|
||||
Set-ItemProperty -Path "HKLM:\$RegPath" -Name "Path" -Value $NewValue -Force | Out-Null
|
||||
$Hive.close()
|
||||
$RegKey.close()
|
||||
|
||||
# Replace Notepad
|
||||
## Currently broken ##
|
||||
# $RegPath = "HKLM:\WinPE-SW\Microsoft\Windows NT\CurrentVersion\Image File Execution Options\notepad.exe"
|
||||
# $NewValue = 'wscript "X:\WK\NotepadPlusPlus\npp.vbs"'
|
||||
# New-Item -Path $RegPath -Force | Out-Null
|
||||
# New-ItemProperty -Path $RegPath -Name "Debugger" -Value $NewValue -Force | Out-Null
|
||||
$RegPath = "HKLM:\WinPE-SW\Microsoft\Windows NT\CurrentVersion\Image File Execution Options\notepad.exe"
|
||||
$NewValue = 'cmd /c "%SystemDrive%\.bin\NotepadPlusPlus\npp.cmd"'
|
||||
New-Item -Path $RegPath -Force | Out-Null
|
||||
New-ItemProperty -Path $RegPath -Name "Debugger" -Value $NewValue -Force | Out-Null
|
||||
|
||||
# Run garbage collection to release potential stale handles
|
||||
## Credit: https://jrich523.wordpress.com/2012/03/06/powershell-loading-and-unloading-registry-hives/
|
||||
|
|
|
|||
3
.gitignore
vendored
3
.gitignore
vendored
|
|
@ -1,11 +1,10 @@
|
|||
**/__pycache__/*
|
||||
*.bak
|
||||
*.iso
|
||||
.bin/Scripts/__pycache__
|
||||
.bin/tmp
|
||||
Drivers
|
||||
Logs
|
||||
Mount
|
||||
PEFiles
|
||||
Scripts/__pycache__
|
||||
WK/amd64/
|
||||
WK/x86/
|
||||
|
|
|
|||
1006
Scripts/functions.py
1006
Scripts/functions.py
File diff suppressed because it is too large
Load diff
153
Scripts/functions/backup.py
Normal file
153
Scripts/functions/backup.py
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
# Wizard Kit PE: Functions - Backup
|
||||
|
||||
from functions.disk import *
|
||||
|
||||
# Regex
|
||||
REGEX_BAD_PATH_NAMES = re.compile(
|
||||
r'([<>:"/\\\|\?\*]'
|
||||
r'|^(CON|PRN|AUX|NUL|COM\d*|LPT\d*)$)'
|
||||
r'|^\s+'
|
||||
r'|[\s\.]+$',
|
||||
re.IGNORECASE)
|
||||
|
||||
def backup_partition(disk, partition):
|
||||
if par['Image Exists'] or par['Number'] in disk['Bad Partitions']:
|
||||
raise GenericAbort
|
||||
|
||||
cmd = [
|
||||
global_vars['Tools']['wimlib-imagex'],
|
||||
'capture'
|
||||
'{}:\\'.format(par['Letter']),
|
||||
par['Image Path'],
|
||||
par['Image Name'], # Image name
|
||||
par['Image Name'], # Image description
|
||||
' --compress=none',
|
||||
]
|
||||
dest_dir = re.sub(r'(.*)\\.*$', r'\1', par['Image Path'], re.IGNORECASE)
|
||||
os.makedirs(dest_dir, exist_ok=True)
|
||||
run_program(cmd)
|
||||
|
||||
def fix_path(path):
|
||||
return REGEX_BAD_PATH_NAMES.sub('_', path)
|
||||
|
||||
def prep_disk_for_backup(destination, disk, ticket_number):
|
||||
disk['Clobber Risk'] = []
|
||||
width = len(str(len(disk['Partitions'])))
|
||||
|
||||
# Get partition totals
|
||||
disk['Bad Partitions'] = [par['Number'] for par in disk['Partitions']
|
||||
if is_bad_partition(partition)]
|
||||
num_valid_partitions = len(disk['Partitions']) - len(disk['Bad Partitions'])
|
||||
disk['Valid Partitions'] = num_valid_partitions
|
||||
if disk['Valid Partitions'] <= 0:
|
||||
print_error('ERROR: No partitions can be backed up for this disk')
|
||||
raise GenericAbort
|
||||
|
||||
# Prep partitions
|
||||
for par in disk['Partitions']:
|
||||
display = 'Partition {num:>{width}}:\t{size} {fs}'.format(
|
||||
num = par['Number'],
|
||||
width = width,
|
||||
size = par['Size'],
|
||||
fs = par['FileSystem'])
|
||||
|
||||
if par['Number'] in disk['Bad Partitions']:
|
||||
# Set display string using partition description & OS type
|
||||
display = ' * {display}\t\t{q}{name}{q}\t{desc} ({os})'.format(
|
||||
display = display,
|
||||
q = '"' if par['Name'] != '' else '',
|
||||
name = par['Name'],
|
||||
desc = par['Description'],
|
||||
os = par['OS'])
|
||||
display = '{YELLOW}{display}{CLEAR}'.format(
|
||||
display=display, **COLORS)
|
||||
else:
|
||||
# Update info for WIM capturing
|
||||
par['Image Name'] = par['Name'] if par['Name'] else 'Unknown'
|
||||
if 'IP' in destination:
|
||||
par['Image Path'] = r'\\{}\{}\{}'.format(
|
||||
destination['IP'], destination['Share'], ticket_number)
|
||||
else:
|
||||
par['Image Path'] = r'{}:\{}'.format(
|
||||
ticket_number, destination['Letter'])
|
||||
par['Image Path'] += r'\{}_{}.wim'.format(
|
||||
par['Number'], par['Image Name'])
|
||||
par['Image Path'] = fix_path(par['Image Path'])
|
||||
|
||||
# Check for existing backups
|
||||
par['Image Exists'] = os.path.exists(par['Image Path'])
|
||||
if par['Image Exists']:
|
||||
disk['Clobber Risk'].append(par['Number'])
|
||||
display = '{} + {}'.format(COLORS['BLUE'], display)
|
||||
else:
|
||||
display = '{} {}'.format(COLORS['CLEAR'], display)
|
||||
|
||||
# Append rest of Display String for valid/clobber partitions
|
||||
display += ' (Used: {used})\t{q}{name}{q}{CLEAR}'.format(
|
||||
used = par['Used Space'],
|
||||
q = '"' if par['Name'] != '' else '',
|
||||
name = par['Name'],
|
||||
**COLORS)
|
||||
# For all partitions
|
||||
par['Display String'] = display
|
||||
|
||||
# Set description for bad partitions
|
||||
warnings = '\n'
|
||||
if disk['Bad Partitions']:
|
||||
warnings += '{} * Unsupported filesystem{}\n'.format(
|
||||
COLORS['YELLOW'], COLORS['CLEAR'])
|
||||
if disk['Clobber Risk']:
|
||||
warnings += '{} + Backup exists on {}{}\n'.format(
|
||||
COLORS['BLUE'], destination['Name'], COLORS['CLEAR'])
|
||||
if disk['Bad Partitions'] or disk['Clobber Risk']:
|
||||
warnings += '\n{}Marked partition(s) will NOT be backed up.{}\n'.format(
|
||||
COLORS['YELLOW'], COLORS['CLEAR'])
|
||||
disk['Backup Warnings'] = warnings
|
||||
|
||||
def select_backup_destination(auto_select=True):
|
||||
# Build menu
|
||||
destinations = [s for s in BACKUP_SERVERS if s['Mounted']]
|
||||
actions = [
|
||||
{'Name': 'Main Menu', 'Letter': 'M'},
|
||||
]
|
||||
|
||||
# Size check
|
||||
for dest in destinations:
|
||||
if 'IP' in dest:
|
||||
dest['Usage'] = shutil.disk_usage(r'\\{IP}\{Share}'.format(**dest))
|
||||
else:
|
||||
dest['Usage'] = shutil.disk_usage('{}:\\'.format(dest['Letter']))
|
||||
dest['Free Space'] = human_readable_size(dest['Usage'].free)
|
||||
dest['Display Name'] = '{Name} ({Free Space} available)'.format(**dest)
|
||||
|
||||
# Bail
|
||||
if not destinations:
|
||||
print_warning('No backup destinations found.')
|
||||
raise GenericAbort
|
||||
|
||||
# Skip menu?
|
||||
if len(destinations) == 1 and auto_select:
|
||||
return destinations[0]
|
||||
|
||||
selection = menu_select(
|
||||
title = 'Where are we backing up to?',
|
||||
main_entries = destinations,
|
||||
action_entries = actions)
|
||||
if selection == 'M':
|
||||
raise GenericAbort
|
||||
else:
|
||||
return destinations[int(selection)-1]
|
||||
|
||||
def verify_wim_backup(partition):
|
||||
if not os.path.exists(partition['Image Path']):
|
||||
raise PathNotFoundError
|
||||
cmd = [
|
||||
global_vars['Tools']['wimlib-imagex'],
|
||||
'verify',
|
||||
partition['Image Path'],
|
||||
' --nocheck',
|
||||
]
|
||||
run_program(cmd)
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("This file is not meant to be called directly.")
|
||||
701
Scripts/functions/common.py
Normal file
701
Scripts/functions/common.py
Normal file
|
|
@ -0,0 +1,701 @@
|
|||
# Wizard Kit PE: Functions - Common
|
||||
|
||||
import os
|
||||
import psutil
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
import winreg
|
||||
|
||||
from subprocess import CalledProcessError
|
||||
|
||||
from settings.main import *
|
||||
from settings.tools import *
|
||||
|
||||
# Global variables
|
||||
global_vars = {}
|
||||
|
||||
# STATIC VARIABLES
|
||||
COLORS = {
|
||||
'CLEAR': '\033[0m',
|
||||
'RED': '\033[31m',
|
||||
'GREEN': '\033[32m',
|
||||
'YELLOW': '\033[33m',
|
||||
'BLUE': '\033[34m'
|
||||
}
|
||||
HKU = winreg.HKEY_USERS
|
||||
HKCU = winreg.HKEY_CURRENT_USER
|
||||
HKLM = winreg.HKEY_LOCAL_MACHINE
|
||||
|
||||
# Error Classes
|
||||
class BIOSKeyNotFoundError(Exception):
|
||||
pass
|
||||
|
||||
class BinNotFoundError(Exception):
|
||||
pass
|
||||
|
||||
class GenericAbort(Exception):
|
||||
pass
|
||||
|
||||
class GenericError(Exception):
|
||||
pass
|
||||
|
||||
class GenericRepair(Exception):
|
||||
pass
|
||||
|
||||
class MultipleInstallationsError(Exception):
|
||||
pass
|
||||
|
||||
class NotInstalledError(Exception):
|
||||
pass
|
||||
|
||||
class NoProfilesError(Exception):
|
||||
pass
|
||||
|
||||
class PathNotFoundError(Exception):
|
||||
pass
|
||||
|
||||
class UnsupportedOSError(Exception):
|
||||
pass
|
||||
|
||||
# General functions
|
||||
def abort():
|
||||
"""Abort script."""
|
||||
print_warning('Aborted.')
|
||||
sleep(5)
|
||||
exit_script()
|
||||
|
||||
def ask(prompt='Kotaero!'):
|
||||
"""Prompt the user with a Y/N question, log answer, and return a bool."""
|
||||
answer = None
|
||||
prompt = '{} [Y/N]: '.format(prompt)
|
||||
while answer is None:
|
||||
tmp = input(prompt)
|
||||
if re.search(r'^y(es|)$', tmp, re.IGNORECASE):
|
||||
answer = True
|
||||
elif re.search(r'^n(o|ope|)$', tmp, re.IGNORECASE):
|
||||
answer = False
|
||||
message = '{prompt}{answer_text}'.format(
|
||||
prompt = prompt,
|
||||
answer_text = 'Yes' if answer else 'No')
|
||||
print_log(message=message)
|
||||
return answer
|
||||
|
||||
def clear_screen():
|
||||
"""Simple wrapper for cls."""
|
||||
os.system('cls')
|
||||
|
||||
def convert_to_bytes(size):
|
||||
"""Convert human-readable size str to bytes and return an int."""
|
||||
size = str(size)
|
||||
tmp = re.search(r'(\d+)\s+([KMGT]B)', size.upper())
|
||||
if tmp:
|
||||
size = int(tmp.group(1))
|
||||
units = tmp.group(2)
|
||||
if units == 'TB':
|
||||
size *= 1099511627776
|
||||
elif units == 'GB':
|
||||
size *= 1073741824
|
||||
elif units == 'MB':
|
||||
size *= 1048576
|
||||
elif units == 'KB':
|
||||
size *= 1024
|
||||
else:
|
||||
return -1
|
||||
|
||||
return size
|
||||
|
||||
def exit_script(return_value=0):
|
||||
"""Exits the script after some cleanup and opens the log (if set)."""
|
||||
# Remove dirs (if empty)
|
||||
for dir in ['BackupDir', 'LogDir', 'TmpDir']:
|
||||
try:
|
||||
dir = global_vars[dir]
|
||||
os.rmdir(dir)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Open Log (if it exists)
|
||||
log = global_vars.get('LogFile', '')
|
||||
if log and os.path.exists(log):
|
||||
try:
|
||||
extract_item('NotepadPlusPlus', silent=True)
|
||||
popen_program(
|
||||
[global_vars['Tools']['NotepadPlusPlus'],
|
||||
global_vars['LogFile']])
|
||||
except Exception:
|
||||
print_error('ERROR: Failed to extract Notepad++ and open log.')
|
||||
pause('Press Enter to exit...')
|
||||
|
||||
# Kill Caffeine if still running
|
||||
kill_process('caffeine.exe')
|
||||
|
||||
# Exit
|
||||
sys.exit(return_value)
|
||||
|
||||
def extract_item(item, filter='', silent=False):
|
||||
"""Extract item from .cbin into .bin."""
|
||||
cmd = [
|
||||
global_vars['Tools']['SevenZip'], 'x', '-aos', '-bso0', '-bse0',
|
||||
'-p{ArchivePassword}'.format(**global_vars),
|
||||
r'-o{BinDir}\{item}'.format(item=item, **global_vars),
|
||||
r'{CBinDir}\{item}.7z'.format(item=item, **global_vars),
|
||||
filter]
|
||||
if not silent:
|
||||
print_standard('Extracting "{item}"...'.format(item=item))
|
||||
try:
|
||||
run_program(cmd)
|
||||
except subprocess.CalledProcessError:
|
||||
if not silent:
|
||||
print_warning('WARNING: Errors encountered while exctracting data')
|
||||
|
||||
def get_ticket_number():
|
||||
"""Get TicketNumber from user, save in LogDir, and return as str."""
|
||||
ticket_number = None
|
||||
while ticket_number is None:
|
||||
_input = input('Enter ticket number: ')
|
||||
if re.match(r'^([0-9]+([-_]?\w+|))$', _input):
|
||||
ticket_number = _input
|
||||
with open(r'{LogDir}\TicketNumber'.format(**global_vars), 'w') as f:
|
||||
f.write(ticket_number)
|
||||
return ticket_number
|
||||
|
||||
def human_readable_size(size, decimals=0):
|
||||
"""Convert size in bytes to a human-readable format and return a str."""
|
||||
# Prep string formatting
|
||||
width = 3+decimals
|
||||
if decimals > 0:
|
||||
width += 1
|
||||
|
||||
# Convert size to int
|
||||
try:
|
||||
size = int(size)
|
||||
except ValueError:
|
||||
size = convert_to_bytes(size)
|
||||
|
||||
# Verify we have a valid size
|
||||
if size < 0:
|
||||
return '{size:>{width}} b'.format(size='???', width=width)
|
||||
|
||||
# Convert to sensible units
|
||||
if size >= 1099511627776:
|
||||
size /= 1099511627776
|
||||
units = 'Tb'
|
||||
elif size >= 1073741824:
|
||||
size /= 1073741824
|
||||
units = 'Gb'
|
||||
elif size >= 1048576:
|
||||
size /= 1048576
|
||||
units = 'Mb'
|
||||
elif size >= 1024:
|
||||
size /= 1024
|
||||
units = 'Kb'
|
||||
else:
|
||||
units = ' b'
|
||||
|
||||
# Return
|
||||
return '{size:>{width}.{decimals}f} {units}'.format(
|
||||
size=size, width=width, decimals=decimals, units=units)
|
||||
|
||||
def kill_process(name):
|
||||
"""Kill any running caffeine.exe processes."""
|
||||
for proc in psutil.process_iter():
|
||||
if proc.name() == name:
|
||||
proc.kill()
|
||||
|
||||
def major_exception():
|
||||
"""Display traceback and exit"""
|
||||
print_error('Major exception')
|
||||
print_warning(SUPPORT_MESSAGE)
|
||||
print(traceback.format_exc())
|
||||
print_log(traceback.format_exc())
|
||||
sleep(30)
|
||||
pause('Press Enter to exit...')
|
||||
exit_script(1)
|
||||
|
||||
def menu_select(title='~ Untitled Menu ~',
|
||||
prompt='Please make a selection', secret_exit=False,
|
||||
main_entries=[], action_entries=[], disabled_label='DISABLED'):
|
||||
"""Display options in a menu and return selected option as a str."""
|
||||
# Bail early
|
||||
if not main_entries and not action_entries:
|
||||
raise Exception("MenuError: No items given")
|
||||
|
||||
# Set title
|
||||
if 'Title' in global_vars:
|
||||
title = '{}\n\n{}'.format(global_vars['Title'])
|
||||
|
||||
# Build menu
|
||||
menu_splash = '{}\n\n'.format(title)
|
||||
width = len(str(len(main_entries)))
|
||||
valid_answers = []
|
||||
if (secret_exit):
|
||||
valid_answers.append('Q')
|
||||
|
||||
# Add main entries
|
||||
for i in range(len(main_entries)):
|
||||
entry = main_entries[i]
|
||||
# Add Spacer
|
||||
if ('CRLF' in entry):
|
||||
menu_splash += '\n'
|
||||
entry_str = '{number:>{width}}: {name}'.format(
|
||||
number = i+1,
|
||||
width = width,
|
||||
name = entry.get('Display Name', entry['Name']))
|
||||
if entry.get('Disabled', False):
|
||||
entry_str = '{YELLOW}{entry_str} ({disabled}){CLEAR}'.format(
|
||||
entry_str = entry_str,
|
||||
disabled = disabled_label,
|
||||
**COLORS)
|
||||
else:
|
||||
valid_answers.append(str(i+1))
|
||||
menu_splash += '{}\n'.format(entry_str)
|
||||
menu_splash += '\n'
|
||||
|
||||
# Add action entries
|
||||
for entry in action_entries:
|
||||
# Add Spacer
|
||||
if ('CRLF' in entry):
|
||||
menu_splash += '\n'
|
||||
valid_answers.append(entry['Letter'])
|
||||
menu_splash += '{letter:>{width}}: {name}\n'.format(
|
||||
letter = entry['Letter'].upper(),
|
||||
width = len(str(len(action_entries))),
|
||||
name = entry['Name'])
|
||||
menu_splash += '\n'
|
||||
|
||||
answer = ''
|
||||
|
||||
while (answer.upper() not in valid_answers):
|
||||
os.system('cls')
|
||||
print(menu_splash)
|
||||
answer = input('{}: '.format(prompt))
|
||||
|
||||
return answer.upper()
|
||||
|
||||
def non_clobber_rename(full_path):
|
||||
"""Append suffix to path, if necessary, to avoid clobbering path"""
|
||||
new_path = full_path
|
||||
_i = 1;
|
||||
while os.path.exists(new_path):
|
||||
new_path = '{path}_{i}'.format(i=_i, path=full_path)
|
||||
_i += 1
|
||||
|
||||
return new_path
|
||||
|
||||
def pause(prompt='Press Enter to continue... '):
|
||||
"""Simple pause implementation."""
|
||||
input(prompt)
|
||||
|
||||
def ping(addr='google.com'):
|
||||
"""Attempt to ping addr."""
|
||||
cmd = ['ping', '-n', '2', addr]
|
||||
run_program(cmd)
|
||||
|
||||
def popen_program(cmd, pipe=False, minimized=False, shell=False, **kwargs):
|
||||
"""Run program and return a subprocess.Popen object."""
|
||||
startupinfo=None
|
||||
if minimized:
|
||||
startupinfo = subprocess.STARTUPINFO()
|
||||
startupinfo.dwFlags = subprocess.STARTF_USESHOWWINDOW
|
||||
startupinfo.wShowWindow = 6
|
||||
|
||||
if pipe:
|
||||
popen_obj = subprocess.Popen(cmd, shell=shell, startupinfo=startupinfo,
|
||||
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
else:
|
||||
popen_obj = subprocess.Popen(cmd, shell=shell, startupinfo=startupinfo)
|
||||
|
||||
return popen_obj
|
||||
|
||||
def print_error(*args, **kwargs):
|
||||
"""Prints message to screen in RED."""
|
||||
print_standard(*args, color=COLORS['RED'], **kwargs)
|
||||
|
||||
def print_info(*args, **kwargs):
|
||||
"""Prints message to screen in BLUE."""
|
||||
print_standard(*args, color=COLORS['BLUE'], **kwargs)
|
||||
|
||||
def print_standard(message='Generic info',
|
||||
color=None, end='\n', timestamp=True, **kwargs):
|
||||
"""Prints message to screen and log (if set)."""
|
||||
display_message = message
|
||||
if color:
|
||||
display_message = color + message + COLORS['CLEAR']
|
||||
# **COLORS is used below to support non-"standard" color printing
|
||||
print(display_message.format(**COLORS), end=end, **kwargs)
|
||||
print_log(message, end, timestamp)
|
||||
|
||||
def print_success(*args, **kwargs):
|
||||
"""Prints message to screen in GREEN."""
|
||||
print_standard(*args, color=COLORS['GREEN'], **kwargs)
|
||||
|
||||
def print_warning(*args, **kwargs):
|
||||
"""Prints message to screen in YELLOW."""
|
||||
print_standard(*args, color=COLORS['YELLOW'], **kwargs)
|
||||
|
||||
def print_log(message='', end='\n', timestamp=True):
|
||||
time_str = time.strftime("%Y-%m-%d %H%M%z: ") if timestamp else ''
|
||||
if 'LogFile' in global_vars and global_vars['LogFile'] is not None:
|
||||
with open(global_vars['LogFile'], 'a') as f:
|
||||
for line in message.splitlines():
|
||||
f.write('{timestamp}{line}{end}'.format(
|
||||
timestamp = time_str,
|
||||
line = line,
|
||||
end = end))
|
||||
|
||||
def run_program(cmd, args=[], check=True, pipe=True, shell=False):
|
||||
"""Run program and return a subprocess.CompletedProcess object."""
|
||||
if args:
|
||||
# Deprecated so let's raise an exception to find & fix all occurances
|
||||
print_error('ERROR: Using args is no longer supported.')
|
||||
raise Exception
|
||||
cmd = [c for c in cmd if c]
|
||||
if shell:
|
||||
cmd = ' '.join(cmd)
|
||||
|
||||
if pipe:
|
||||
process_return = subprocess.run(cmd, check=check, shell=shell,
|
||||
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
else:
|
||||
process_return = subprocess.run(cmd, check=check, shell=shell)
|
||||
|
||||
return process_return
|
||||
|
||||
def set_title(title='~Some Title~'):
|
||||
"""Set title.
|
||||
|
||||
Used for window title and menu titles."""
|
||||
global_vars['Title'] = title
|
||||
os.system('title {}'.format(title))
|
||||
|
||||
def show_info(message='~Some message~', info='~Some info~', indent=8, width=32):
|
||||
"""Display info with formatting."""
|
||||
print_standard('{indent}{message:<{width}}{info}'.format(
|
||||
indent=' '*indent, width=width, message=message, info=info))
|
||||
|
||||
def sleep(seconds=2):
|
||||
"""Wait for a while."""
|
||||
time.sleep(seconds)
|
||||
|
||||
def stay_awake():
|
||||
"""Prevent the system from sleeping or hibernating."""
|
||||
# Bail if caffeine is already running
|
||||
for proc in psutil.process_iter():
|
||||
if proc.name() == 'caffeine.exe':
|
||||
return
|
||||
# Extract and run
|
||||
extract_item('Caffeine', silent=True)
|
||||
try:
|
||||
popen_program(global_vars['Tools']['Caffeine'])
|
||||
except Exception:
|
||||
print_error('ERROR: No caffeine available.')
|
||||
print_warning('Please set the power setting to High Performance.')
|
||||
|
||||
def get_exception(s):
|
||||
"""Get exception by name, returns Exception object."""
|
||||
return getattr(sys.modules[__name__], s)
|
||||
|
||||
def try_and_print(message='Trying...',
|
||||
function=None, cs='CS', ns='NS', other_results={},
|
||||
catch_all=True, print_return=False, silent_function=True,
|
||||
indent=8, width=32, *args, **kwargs):
|
||||
"""Run function, print if successful or not, and return dict.
|
||||
|
||||
other_results is in the form of
|
||||
{
|
||||
'Warning': {'ExceptionClassName': 'Result Message'},
|
||||
'Error': {'ExceptionClassName': 'Result Message'}
|
||||
}
|
||||
The the ExceptionClassNames will be excepted conditions
|
||||
and the result string will be printed in the correct color.
|
||||
catch_all=False will result in unspecified exceptions being re-raised."""
|
||||
err = None
|
||||
out = None
|
||||
w_exceptions = other_results.get('Warning', {}).keys()
|
||||
w_exceptions = tuple(get_exception(e) for e in w_exceptions)
|
||||
e_exceptions = other_results.get('Error', {}).keys()
|
||||
e_exceptions = tuple(get_exception(e) for e in e_exceptions)
|
||||
w_results = other_results.get('Warning', {})
|
||||
e_results = other_results.get('Error', {})
|
||||
|
||||
# Run function and catch errors
|
||||
print_standard('{indent}{message:<{width}}'.format(
|
||||
indent=' '*indent, message=message, width=width), end='', flush=True)
|
||||
try:
|
||||
out = function(*args, **kwargs)
|
||||
if print_return:
|
||||
print_standard(out[0], timestamp=False)
|
||||
for item in out[1:]:
|
||||
print_standard('{indent}{item}'.format(
|
||||
indent=' '*(indent+width), item=item))
|
||||
elif silent_function:
|
||||
print_success(cs, timestamp=False)
|
||||
except w_exceptions as e:
|
||||
_result = w_results.get(e.__class__.__name__, 'Warning')
|
||||
print_warning(_result, timestamp=False)
|
||||
err = e
|
||||
except e_exceptions as e:
|
||||
_result = e_results.get(e.__class__.__name__, 'Error')
|
||||
print_error(_result, timestamp=False)
|
||||
err = e
|
||||
except Exception:
|
||||
print_error(ns, timestamp=False)
|
||||
err = traceback.format_exc()
|
||||
|
||||
# Return or raise?
|
||||
if err and not catch_all:
|
||||
raise
|
||||
else:
|
||||
return {'CS': not bool(err), 'Error': err, 'Out': out}
|
||||
|
||||
def upload_data(path, file):
|
||||
"""Add CLIENT_INFO_SERVER to authorized connections and upload file."""
|
||||
if not ENABLED_UPLOAD_DATA:
|
||||
raise GenericError('Feature disabled.')
|
||||
|
||||
extract_item('PuTTY', filter='wizkit.ppk psftp.exe', silent=True)
|
||||
|
||||
# Authorize connection to the server
|
||||
winreg.CreateKey(HKCU, r'Software\SimonTatham\PuTTY\SshHostKeys')
|
||||
with winreg.OpenKey(HKCU, r'Software\SimonTatham\PuTTY\SshHostKeys',
|
||||
access=winreg.KEY_WRITE) as key:
|
||||
winreg.SetValueEx(key,
|
||||
'rsa2@22:{IP}'.format(**CLIENT_INFO_SERVER), 0,
|
||||
winreg.REG_SZ, CLIENT_INFO_SERVER['RegEntry'])
|
||||
|
||||
# Write batch file
|
||||
with open(r'{TmpDir}\psftp.batch'.format(**global_vars),
|
||||
'w', encoding='ascii') as f:
|
||||
f.write('lcd "{path}"\n'.format(path=path))
|
||||
f.write('cd "{Share}"\n'.format(**CLIENT_INFO_SERVER))
|
||||
f.write('mkdir {TicketNumber}\n'.format(**global_vars))
|
||||
f.write('cd {TicketNumber}\n'.format(**global_vars))
|
||||
f.write('put "{file}"\n'.format(file=file))
|
||||
|
||||
# Upload Info
|
||||
cmd = [
|
||||
global_vars['Tools']['PuTTY-PSFTP'],
|
||||
'-noagent',
|
||||
'-i', r'{BinDir}\PuTTY\wizkit.ppk'.format(**global_vars),
|
||||
'{User}@{IP}'.format(**CLIENT_INFO_SERVER),
|
||||
'-b', r'{TmpDir}\psftp.batch'.format(**global_vars)]
|
||||
run_program(cmd)
|
||||
|
||||
def upload_info():
|
||||
"""Upload compressed Info file to the NAS as set in settings.main.py."""
|
||||
if not ENABLED_UPLOAD_DATA:
|
||||
raise GenericError('Feature disabled.')
|
||||
|
||||
path = '{ClientDir}'.format(**global_vars)
|
||||
file = 'Info_{Date-Time}.7z'.format(**global_vars)
|
||||
upload_data(path, file)
|
||||
|
||||
def compress_info():
|
||||
"""Compress ClientDir info folders with 7-Zip for upload_info()."""
|
||||
path = '{ClientDir}'.format(**global_vars)
|
||||
file = 'Info_{Date-Time}.7z'.format(**global_vars)
|
||||
_cmd = [
|
||||
global_vars['Tools']['SevenZip'],
|
||||
'a', '-t7z', '-mx=9', '-bso0', '-bse0',
|
||||
r'{}\{}'.format(path, file),
|
||||
r'{ClientDir}\Info'.format(**global_vars)]
|
||||
run_program(_cmd)
|
||||
|
||||
def wait_for_process(name, poll_rate=3):
|
||||
"""Wait for process by name."""
|
||||
running = True
|
||||
while running:
|
||||
sleep(poll_rate)
|
||||
running = False
|
||||
for proc in psutil.process_iter():
|
||||
if re.search(r'^{}'.format(name), proc.name(), re.IGNORECASE):
|
||||
running = True
|
||||
sleep(1)
|
||||
|
||||
# global_vars functions
|
||||
def init_global_vars():
|
||||
"""Sets global variables based on system info."""
|
||||
print_info('Initializing')
|
||||
os.system('title Wizard Kit')
|
||||
init_functions = [
|
||||
['Checking .bin...', find_bin],
|
||||
['Checking environment...', set_common_vars],
|
||||
['Checking OS...', check_os],
|
||||
['Checking tools...', check_tools],
|
||||
['Creating folders...', make_tmp_dirs],
|
||||
['Clearing collisions...', clean_env_vars],
|
||||
]
|
||||
try:
|
||||
for f in init_functions:
|
||||
try_and_print(
|
||||
message=f[0], function=f[1],
|
||||
cs='Done', ns='Error', catch_all=False)
|
||||
except:
|
||||
major_exception()
|
||||
|
||||
def check_os():
|
||||
"""Set OS specific variables."""
|
||||
tmp = {}
|
||||
|
||||
# Query registry
|
||||
_reg_path = winreg.OpenKey(
|
||||
HKLM, r'SOFTWARE\Microsoft\Windows NT\CurrentVersion')
|
||||
for key in ['CSDVersion', 'CurrentBuild', 'CurrentBuildNumber',
|
||||
'CurrentVersion', 'ProductName']:
|
||||
try:
|
||||
tmp[key] = winreg.QueryValueEx(_reg_path, key)[0]
|
||||
if key in ['CurrentBuild', 'CurrentBuildNumber']:
|
||||
tmp[key] = int(tmp[key])
|
||||
except ValueError:
|
||||
# Couldn't convert Build to int so this should be interesting...
|
||||
tmp[key] = 0
|
||||
except Exception:
|
||||
tmp[key] = 'Unknown'
|
||||
|
||||
# Determine OS bit depth
|
||||
tmp['Arch'] = 32
|
||||
if 'PROGRAMFILES(X86)' in global_vars['Env']:
|
||||
tmp['Arch'] = 64
|
||||
|
||||
# Determine OS Name
|
||||
tmp['Name'] = '{ProductName} {CSDVersion}'.format(**tmp)
|
||||
if tmp['CurrentBuild'] == 9600:
|
||||
tmp['Name'] += ' Update' # Win 8.1u
|
||||
if tmp['CurrentBuild'] == 10240:
|
||||
tmp['Name'] += ' Release 1507 "Threshold 1"'
|
||||
if tmp['CurrentBuild'] == 10586:
|
||||
tmp['Name'] += ' Release 1511 "Threshold 2"'
|
||||
if tmp['CurrentBuild'] == 14393:
|
||||
tmp['Name'] += ' Release 1607 "Redstone 1" / "Anniversary Update"'
|
||||
if tmp['CurrentBuild'] == 15063:
|
||||
tmp['Name'] += ' Release 1703 "Redstone 2" / "Creators Update"'
|
||||
if tmp['CurrentBuild'] == 16299:
|
||||
tmp['Name'] += ' Release 1709 "Redstone 3" / "Fall Creators Update"'
|
||||
tmp['Name'] = tmp['Name'].replace('Service Pack ', 'SP')
|
||||
tmp['Name'] = tmp['Name'].replace('Unknown Release', 'Release')
|
||||
tmp['Name'] = re.sub(r'\s+', ' ', tmp['Name'])
|
||||
|
||||
# Determine OS version
|
||||
name = '{Name} x{Arch}'.format(**tmp)
|
||||
if tmp['CurrentVersion'] == '6.0':
|
||||
tmp['Version'] = 'Vista'
|
||||
name += ' (very outdated)'
|
||||
elif tmp['CurrentVersion'] == '6.1':
|
||||
tmp['Version'] = '7'
|
||||
if tmp['CSDVersion'] == 'Service Pack 1':
|
||||
name += ' (outdated)'
|
||||
else:
|
||||
name += ' (very outdated)'
|
||||
elif tmp['CurrentVersion'] in ['6.2', '6.3']:
|
||||
if int(tmp['CurrentBuildNumber']) <= 9600:
|
||||
tmp['Version'] = '8'
|
||||
elif int(tmp['CurrentBuildNumber']) >= 10240:
|
||||
tmp['Version'] = '10'
|
||||
if tmp['CurrentBuild'] in [9200, 10240, 10586]:
|
||||
name += ' (very outdated)'
|
||||
elif tmp['CurrentBuild'] in [9600, 14393, 15063]:
|
||||
name += ' (outdated)'
|
||||
elif tmp['CurrentBuild'] == 16299:
|
||||
pass # Current Win10
|
||||
else:
|
||||
name += ' (unrecognized)'
|
||||
tmp['DisplayName'] = name
|
||||
|
||||
# == vista ==
|
||||
# 6.0.6000
|
||||
# 6.0.6001
|
||||
# 6.0.6002
|
||||
# ==== 7 ====
|
||||
# 6.1.7600
|
||||
# 6.1.7601
|
||||
# 6.1.7602
|
||||
# ==== 8 ====
|
||||
# 6.2.9200
|
||||
# === 8.1 ===
|
||||
# 6.3.9200
|
||||
# === 8.1u ==
|
||||
# 6.3.9600
|
||||
# === 10 v1507 "Threshold 1" ==
|
||||
# 6.3.10240
|
||||
# === 10 v1511 "Threshold 2" ==
|
||||
# 6.3.10586
|
||||
# === 10 v1607 "Redstone 1" "Anniversary Update" ==
|
||||
# 6.3.14393
|
||||
# === 10 v1703 "Redstone 2" "Creators Update" ==
|
||||
# 6.3.15063
|
||||
# === 10 v1709 "Redstone 3" "Fall Creators Update" ==
|
||||
# 6.3.16299
|
||||
global_vars['OS'] = tmp
|
||||
|
||||
def check_tools():
|
||||
"""Set tool variables based on OS bit-depth and tool availability."""
|
||||
if global_vars['OS'].get('Arch', 32) == 64:
|
||||
global_vars['Tools'] = {
|
||||
k: v.get('64', v.get('32')) for (k, v) in TOOLS.items()}
|
||||
else:
|
||||
global_vars['Tools'] = {k: v.get('32') for (k, v) in TOOLS.items()}
|
||||
|
||||
# Fix paths
|
||||
global_vars['Tools'] = {k: os.path.join(global_vars['BinDir'], v)
|
||||
for (k, v) in global_vars['Tools'].items()}
|
||||
|
||||
def clean_env_vars():
|
||||
"""Remove conflicting global_vars and env variables.
|
||||
|
||||
This fixes an issue where both global_vars and
|
||||
global_vars['Env'] are expanded at the same time."""
|
||||
for key in global_vars.keys():
|
||||
global_vars['Env'].pop(key, None)
|
||||
|
||||
def find_bin():
|
||||
"""Find .bin folder in the cwd or it's parents."""
|
||||
wd = os.getcwd()
|
||||
base = None
|
||||
while base is None:
|
||||
if os.path.exists('.bin'):
|
||||
base = os.getcwd()
|
||||
break
|
||||
if re.fullmatch(r'\w:\\', os.getcwd()):
|
||||
break
|
||||
os.chdir('..')
|
||||
os.chdir(wd)
|
||||
if base is None:
|
||||
raise BinNotFoundError
|
||||
global_vars['BaseDir'] = base
|
||||
|
||||
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(global_vars['TmpDir'], exist_ok=True)
|
||||
|
||||
def set_common_vars():
|
||||
"""Set common variables."""
|
||||
global_vars['Date'] = time.strftime("%Y-%m-%d")
|
||||
global_vars['Date-Time'] = time.strftime("%Y-%m-%d_%H%M_%z")
|
||||
global_vars['Env'] = os.environ.copy()
|
||||
|
||||
global_vars['ArchivePassword'] = ARCHIVE_PASSWORD
|
||||
global_vars['BinDir'] = r'{BaseDir}\.bin'.format(
|
||||
**global_vars)
|
||||
global_vars['CBinDir'] = r'{BaseDir}\.cbin'.format(
|
||||
**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)
|
||||
global_vars['LogDir'] = r'{ClientDir}\Info\{Date}'.format(
|
||||
**global_vars)
|
||||
global_vars['ProgBackupDir'] = r'{ClientDir}\Backups'.format(
|
||||
**global_vars)
|
||||
global_vars['QuarantineDir'] = r'{ClientDir}\Quarantine'.format(
|
||||
**global_vars)
|
||||
global_vars['TmpDir'] = r'{BinDir}\tmp'.format(
|
||||
**global_vars)
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("This file is not meant to be called directly.")
|
||||
607
Scripts/functions/data.py
Normal file
607
Scripts/functions/data.py
Normal file
|
|
@ -0,0 +1,607 @@
|
|||
# Wizard Kit PE: Functions - Data
|
||||
|
||||
import ctypes
|
||||
|
||||
from operator import itemgetter
|
||||
|
||||
from functions.common import *
|
||||
|
||||
# Classes
|
||||
class LocalDisk():
|
||||
def __init__(self, disk):
|
||||
self.disk = disk
|
||||
self.name = disk.mountpoint.upper()
|
||||
self.path = self.name
|
||||
def is_dir(self):
|
||||
# Should always be true
|
||||
return True
|
||||
|
||||
# Regex
|
||||
REGEX_EXCL_ITEMS = re.compile(
|
||||
r'^(\.(AppleDB|AppleDesktop|AppleDouble'
|
||||
r'|com\.apple\.timemachine\.supported|dbfseventsd'
|
||||
r'|DocumentRevisions-V100.*|DS_Store|fseventsd|PKInstallSandboxManager'
|
||||
r'|Spotlight.*|SymAV.*|symSchedScanLockxz|TemporaryItems|Trash.*'
|
||||
r'|vol|VolumeIcon\.icns)|desktop\.(ini|.*DB|.*DF)'
|
||||
r'|(hiberfil|pagefile)\.sys|lost\+found|Network\.*Trash\.*Folder'
|
||||
r'|Recycle[dr]|System\.*Volume\.*Information|Temporary\.*Items'
|
||||
r'|Thumbs\.db)$',
|
||||
re.IGNORECASE)
|
||||
REGEX_EXCL_ROOT_ITEMS = re.compile(
|
||||
r'^\\?(boot(mgr|nxt)$|Config.msi'
|
||||
r'|(eula|globdata|install|vc_?red)'
|
||||
r'|.*.sys$|System Volume Information|RECYCLER?|\$Recycle\.bin'
|
||||
r'|\$?Win(dows(.old.*|\.~BT|)$|RE_)|\$GetCurrent|Windows10Upgrade'
|
||||
r'|PerfLogs|Program Files|SYSTEM.SAV'
|
||||
r'|.*\.(esd|swm|wim|dd|map|dmg|image)$)',
|
||||
re.IGNORECASE)
|
||||
REGEX_INCL_ROOT_ITEMS = re.compile(
|
||||
r'^\\?(AdwCleaner|(My\s*|)(Doc(uments?( and Settings|)|s?)|Downloads'
|
||||
r'|Media|Music|Pic(ture|)s?|Vid(eo|)s?)'
|
||||
r'|{prefix}(-?Info|-?Transfer|)'
|
||||
r'|(ProgramData|Recovery|Temp.*|Users)$'
|
||||
r'|.*\.(log|txt|rtf|qb\w*|avi|m4a|m4v|mp4|mkv|jpg|png|tiff?)$)'
|
||||
r''.format(prefix=KIT_NAME_SHORT),
|
||||
re.IGNORECASE)
|
||||
REGEX_WIM_FILE = re.compile(
|
||||
r'\.wim$',
|
||||
re.IGNORECASE)
|
||||
REGEX_WINDOWS_OLD = re.compile(
|
||||
r'^\\Win(dows|)\.old',
|
||||
re.IGNORECASE)
|
||||
|
||||
# STATIC VARIABLES
|
||||
FAST_COPY_EXCLUDES = [
|
||||
r'\*.esd',
|
||||
r'\*.swm',
|
||||
r'\*.wim',
|
||||
r'\*.dd',
|
||||
r'\*.dd.tgz',
|
||||
r'\*.dd.txz',
|
||||
r'\*.map',
|
||||
r'\*.dmg',
|
||||
r'\*.image',
|
||||
r'$RECYCLE.BIN',
|
||||
r'$Recycle.Bin',
|
||||
r'.AppleDB',
|
||||
r'.AppleDesktop',
|
||||
r'.AppleDouble',
|
||||
r'.com.apple.timemachine.supported',
|
||||
r'.dbfseventsd',
|
||||
r'.DocumentRevisions-V100*',
|
||||
r'.DS_Store',
|
||||
r'.fseventsd',
|
||||
r'.PKInstallSandboxManager',
|
||||
r'.Spotlight*',
|
||||
r'.SymAV*',
|
||||
r'.symSchedScanLockxz',
|
||||
r'.TemporaryItems',
|
||||
r'.Trash*',
|
||||
r'.vol',
|
||||
r'.VolumeIcon.icns',
|
||||
r'desktop.ini',
|
||||
r'Desktop?DB',
|
||||
r'Desktop?DF',
|
||||
r'hiberfil.sys',
|
||||
r'lost+found',
|
||||
r'Network?Trash?Folder',
|
||||
r'pagefile.sys',
|
||||
r'Recycled',
|
||||
r'RECYCLER',
|
||||
r'System?Volume?Information',
|
||||
r'Temporary?Items',
|
||||
r'Thumbs.db',
|
||||
]
|
||||
FAST_COPY_ARGS = [
|
||||
'/cmd=noexist_only',
|
||||
'/utf8',
|
||||
'/skip_empty_dir',
|
||||
'/linkdest',
|
||||
'/no_ui',
|
||||
'/auto_close',
|
||||
'/exclude={}'.format(';'.join(FAST_COPY_EXCLUDES)),
|
||||
]
|
||||
# Code borrowed from: https://stackoverflow.com/a/29075319
|
||||
SEM_NORMAL = ctypes.c_uint()
|
||||
SEM_FAILCRITICALERRORS = 1
|
||||
SEM_NOOPENFILEERRORBOX = 0x8000
|
||||
SEM_FAIL = SEM_NOOPENFILEERRORBOX | SEM_FAILCRITICALERRORS
|
||||
|
||||
def cleanup_transfer(dest_path):
|
||||
"""Fix attributes and move extraneous items outside the Transfer folder."""
|
||||
try:
|
||||
# Remove dest_path if empty
|
||||
os.rmdir(dest_path)
|
||||
except OSError:
|
||||
pass
|
||||
if not os.path.exists(dest_path):
|
||||
# Bail if dest_path was empty and removed
|
||||
raise Exception
|
||||
|
||||
# Fix attributes
|
||||
cmd = ['attrib', '-a', '-h', '-r', '-s', dest_path]
|
||||
run_program(cmd, check=False)
|
||||
|
||||
for root, dirs, files in os.walk(dest_path, topdown=False):
|
||||
for name in dirs:
|
||||
# Remove empty directories and junction points
|
||||
try:
|
||||
os.rmdir(os.path.join(root, name))
|
||||
except OSError:
|
||||
pass
|
||||
for name in files:
|
||||
# "Remove" files based on exclusion regex
|
||||
if REGEX_EXCL_ITEMS.search(name):
|
||||
# Make dest folder
|
||||
dest_name = root.replace(dest_path, dest_path+'.Removed')
|
||||
os.makedirs(dest_name, exist_ok=True)
|
||||
# Set dest filename
|
||||
dest_name = os.path.join(dest_name, name)
|
||||
dest_name = non_clobber_rename(dest_name)
|
||||
source_name = os.path.join(root, name)
|
||||
try:
|
||||
shutil.move(source_name, dest_name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def is_valid_wim_file(item):
|
||||
"""Checks if the provided os.DirEntry is a valid WIM file, returns bool."""
|
||||
valid = bool(item.is_file() and REGEX_WIM_FILE.search(item.name))
|
||||
if valid:
|
||||
extract_item('wimlib', silent=True)
|
||||
cmd = [global_vars['Tools']['wimlib-imagex'], 'info', item.path]
|
||||
try:
|
||||
run_program(cmd)
|
||||
except subprocess.CalledProcessError:
|
||||
valid = False
|
||||
print_log('WARNING: Image "{}" damaged.'.format(item.name))
|
||||
return valid
|
||||
|
||||
def mount_backup_shares():
|
||||
"""Mount the backup shares unless labeled as already mounted."""
|
||||
for server in BACKUP_SERVERS:
|
||||
# Blindly skip if we mounted earlier
|
||||
if server['Mounted']:
|
||||
continue
|
||||
|
||||
mount_network_share(server)
|
||||
|
||||
def mount_network_share(server):
|
||||
"""Mount a network share defined by server."""
|
||||
# Test connection
|
||||
try:
|
||||
ping(server['IP'])
|
||||
except subprocess.CalledProcessError:
|
||||
print_error(
|
||||
r'Failed to mount \\{Name}\{Share}, {IP} unreachable.'.format(
|
||||
**server))
|
||||
sleep(1)
|
||||
return False
|
||||
|
||||
# Mount
|
||||
cmd = r'net use \\{IP}\{Share} /user:{User} {Pass}'.format(**server)
|
||||
cmd = cmd.split(' ')
|
||||
try:
|
||||
run_program(cmd)
|
||||
except Exception:
|
||||
print_warning(r'Failed to mount \\{Name}\{Share} ({IP})'.format(
|
||||
**server))
|
||||
sleep(1)
|
||||
else:
|
||||
print_info('Mounted {Name}'.format(**server))
|
||||
server['Mounted'] = True
|
||||
|
||||
def run_fast_copy(items, dest):
|
||||
"""Copy items to dest using FastCopy."""
|
||||
if not items:
|
||||
raise Exception
|
||||
|
||||
cmd = [global_vars['Tools']['FastCopy'], *FAST_COPY_ARGS]
|
||||
if 'LogFile' in global_vars:
|
||||
cmd.append('/logfile={LogFile}'.format(**global_vars))
|
||||
cmd.extend(items)
|
||||
cmd.append('/to={}\\'.format(dest))
|
||||
|
||||
run_program(cmd)
|
||||
|
||||
def run_wimextract(source, items, dest):
|
||||
"""Extract items from source WIM to dest folder."""
|
||||
if not items:
|
||||
raise Exception
|
||||
extract_item('wimlib', silent=True)
|
||||
|
||||
# Write files.txt
|
||||
with open(r'{TmpDir}\wim_files.txt'.format(**global_vars), 'w') as f:
|
||||
# Defaults
|
||||
for item in items:
|
||||
f.write('{item}\n'.format(item=item))
|
||||
sleep(1) # For safety?
|
||||
|
||||
# Extract files
|
||||
cmd = [
|
||||
global_vars['Tools']['wimlib-imagex'],
|
||||
'extract',
|
||||
source, '1',
|
||||
r'@{TmpDir}\wim_files.txt'.format(**global_vars),
|
||||
'--dest-dir={}\\'.format(dest),
|
||||
'--no-acls',
|
||||
'--nullglob']
|
||||
run_program(cmd)
|
||||
|
||||
def scan_source(source_obj, dest_path):
|
||||
"""Scan source for files/folders to transfer."""
|
||||
selected_items = []
|
||||
|
||||
if source_obj.is_dir():
|
||||
# File-Based
|
||||
print_standard('Scanning source (folder): {}'.format(source_obj.path))
|
||||
selected_items = scan_source_path(source_obj.path, dest_path)
|
||||
else:
|
||||
# Image-Based
|
||||
if REGEX_WIM_FILE.search(source_obj.name):
|
||||
print_standard('Scanning source (image): {}'.format(
|
||||
source_obj.path))
|
||||
selected_items = scan_source_wim(source_obj.path, dest_path)
|
||||
else:
|
||||
print_error('ERROR: Unsupported image: {}'.format(
|
||||
source_obj.path))
|
||||
raise GenericError
|
||||
|
||||
return selected_items
|
||||
|
||||
def scan_source_path(source_path, dest_path, rel_path=None, interactive=True):
|
||||
"""Scan source folder for files/folders to transfer, returns list.
|
||||
|
||||
This will scan the root and (recursively) any Windows.old folders."""
|
||||
rel_path = '\\' + rel_path if rel_path else ''
|
||||
if rel_path:
|
||||
dest_path = dest_path + rel_path
|
||||
selected_items = []
|
||||
win_olds = []
|
||||
|
||||
# Root items
|
||||
root_items = []
|
||||
for item in os.scandir(source_path):
|
||||
if REGEX_INCL_ROOT_ITEMS.search(item.name):
|
||||
root_items.append(item.path)
|
||||
elif not REGEX_EXCL_ROOT_ITEMS.search(item.name):
|
||||
if (not interactive
|
||||
or ask('Copy: "{}{}" ?'.format(rel_path, item.name))):
|
||||
root_items.append(item.path)
|
||||
if REGEX_WINDOWS_OLD.search(item.name):
|
||||
win_olds.append(item)
|
||||
if root_items:
|
||||
selected_items.append({
|
||||
'Message': '{}Root Items...'.format(rel_path),
|
||||
'Items': root_items.copy(),
|
||||
'Destination': dest_path})
|
||||
|
||||
# Fonts
|
||||
if os.path.exists(r'{}\Windows\Fonts'.format(source_path)):
|
||||
selected_items.append({
|
||||
'Message': '{}Fonts...'.format(rel_path),
|
||||
'Items': [r'{}\Windows\Fonts'.format(rel_path)],
|
||||
'Destination': r'{}\Windows'.format(dest_path)})
|
||||
|
||||
# Registry
|
||||
registry_items = []
|
||||
for folder in ['config', 'OEM']:
|
||||
folder = r'Windows\System32\{}'.format(folder)
|
||||
folder = os.path.join(source_path, folder)
|
||||
if os.path.exists(folder):
|
||||
registry_items.append(folder)
|
||||
if registry_items:
|
||||
selected_items.append({
|
||||
'Message': '{}Registry...'.format(rel_path),
|
||||
'Items': registry_items.copy(),
|
||||
'Destination': r'{}\Windows\System32'.format(dest_path)})
|
||||
|
||||
# Windows.old(s)
|
||||
for old in win_olds:
|
||||
selected_items.append(
|
||||
scan_source_path(
|
||||
old.path, dest_path, rel_path=old.name, interactive=False))
|
||||
|
||||
# Done
|
||||
return selected_items
|
||||
|
||||
def scan_source_wim(source_wim, dest_path, rel_path=None, interactive=True):
|
||||
"""Scan source WIM file for files/folders to transfer, returns list.
|
||||
|
||||
This will scan the root and (recursively) any Windows.old folders."""
|
||||
rel_path = '\\' + rel_path if rel_path else ''
|
||||
selected_items = []
|
||||
win_olds = []
|
||||
|
||||
# Scan source
|
||||
extract_item('wimlib', silent=True)
|
||||
cmd = [
|
||||
global_vars['Tools']['wimlib-imagex'], 'dir',
|
||||
source_wim, '1']
|
||||
try:
|
||||
file_list = run_program(cmd)
|
||||
except subprocess.CalledProcessError:
|
||||
print_error('ERROR: Failed to get file list.')
|
||||
raise
|
||||
|
||||
# Root Items
|
||||
file_list = [i.strip()
|
||||
for i in file_list.stdout.decode('utf-8', 'ignore').splitlines()
|
||||
if i.count('\\') == 1 and i.strip() != '\\']
|
||||
root_items = []
|
||||
if rel_path:
|
||||
file_list = [i.replace(rel_path, '') for i in file_list]
|
||||
for item in file_list:
|
||||
if REGEX_INCL_ROOT_ITEMS.search(item):
|
||||
root_items.append(item)
|
||||
elif not REGEX_EXCL_ROOT_ITEMS.search(item):
|
||||
if (not interactive
|
||||
or ask('Extract: "{}{}" ?'.format(rel_path, item))):
|
||||
root_items.append('{}{}'.format(rel_path, item))
|
||||
if REGEX_WINDOWS_OLD.search(item):
|
||||
win_olds.append(item)
|
||||
if root_items:
|
||||
selected_items.append({
|
||||
'Message': '{}Root Items...'.format(rel_path),
|
||||
'Items': root_items.copy(),
|
||||
'Destination': dest_path})
|
||||
|
||||
# Fonts
|
||||
if wim_contains(source_wim, r'{}Windows\Fonts'.format(rel_path)):
|
||||
selected_items.append({
|
||||
'Message': '{}Fonts...'.format(rel_path),
|
||||
'Items': [r'{}\Windows\Fonts'.format(rel_path)],
|
||||
'Destination': dest_path})
|
||||
|
||||
# Registry
|
||||
registry_items = []
|
||||
for folder in ['config', 'OEM']:
|
||||
folder = r'{}Windows\System32\{}'.format(rel_path, folder)
|
||||
if wim_contains(source_wim, folder):
|
||||
registry_items.append(folder)
|
||||
if registry_items:
|
||||
selected_items.append({
|
||||
'Message': '{}Registry...'.format(rel_path),
|
||||
'Items': registry_items.copy(),
|
||||
'Destination': dest_path})
|
||||
|
||||
# Windows.old(s)
|
||||
for old in win_olds:
|
||||
scan_source_wim(source_wim, dest_path, rel_path=old, interactive=False)
|
||||
|
||||
# Done
|
||||
return selected_items
|
||||
|
||||
def select_destination(folder_path, prompt='Select destination'):
|
||||
"""Select destination drive, returns path as string."""
|
||||
disk = select_volume(prompt)
|
||||
if 'fixed' not in disk['Disk'].opts:
|
||||
folder_path = folder_path.replace('\\', '-')
|
||||
path = '{disk}{folder_path}_{Date}'.format(
|
||||
disk = disk['Disk'].mountpoint,
|
||||
folder_path = folder_path,
|
||||
**global_vars)
|
||||
|
||||
# Avoid merging with existing folder
|
||||
path = non_clobber_rename(path)
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
return path
|
||||
|
||||
def select_source(ticket_number):
|
||||
"""Select backup from those found on the BACKUP_SERVERS for the ticket."""
|
||||
selected_source = None
|
||||
sources = []
|
||||
mount_backup_shares()
|
||||
|
||||
# Check for ticket folders on servers
|
||||
for server in BACKUP_SERVERS:
|
||||
if server['Mounted']:
|
||||
print_standard('Scanning {}...'.format(server['Name']))
|
||||
for d in os.scandir(r'\\{IP}\{Share}'.format(**server)):
|
||||
if (d.is_dir()
|
||||
and d.name.lower().startswith(ticket_number.lower())):
|
||||
# Add folder to sources
|
||||
sources.append({
|
||||
'Name': '{:9}| File-Based: [DIR] {}'.format(
|
||||
server['Name'], d.name),
|
||||
'Server': server,
|
||||
'Source': d})
|
||||
|
||||
# Check for images and subfolders
|
||||
for ticket_path in sources.copy():
|
||||
for item in os.scandir(ticket_path['Source'].path):
|
||||
if item.is_dir():
|
||||
# Add folder to sources
|
||||
sources.append({
|
||||
'Name': r'{:9}| File-Based: [DIR] {}\{}'.format(
|
||||
ticket_path['Server']['Name'], # Server
|
||||
ticket_path['Source'].name, # Ticket folder
|
||||
item.name, # Sub-folder
|
||||
),
|
||||
'Server': ticket_path['Server'],
|
||||
'Source': item})
|
||||
|
||||
# Check for images in folder
|
||||
for subitem in os.scandir(item.path):
|
||||
if REGEX_WIM_FILE.search(item.name):
|
||||
# Add image to sources
|
||||
try:
|
||||
size = human_readable_size(item.stat().st_size)
|
||||
except Exception:
|
||||
size = ' ? ?' # unknown
|
||||
sources.append({
|
||||
'Disabled': bool(not is_valid_wim_file(subitem)),
|
||||
'Name': r'{:9}| Image-Based: {:>7} {}\{}\{}'.format(
|
||||
ticket_path['Server']['Name'], # Server
|
||||
size, # Size (duh)
|
||||
ticket_path['Source'].name, # Ticket folder
|
||||
item.name, # Sub-folder
|
||||
subitem.name, # Image file
|
||||
),
|
||||
'Server': ticket_path['Server'],
|
||||
'Source': subitem})
|
||||
elif REGEX_WIM_FILE.search(item.name):
|
||||
# Add image to sources
|
||||
try:
|
||||
size = human_readable_size(item.stat().st_size)
|
||||
except Exception:
|
||||
size = ' ? ?' # unknown
|
||||
sources.append({
|
||||
'Disabled': bool(not is_valid_wim_file(item)),
|
||||
'Name': r'{:9}| Image-Based: {:>7} {}\{}'.format(
|
||||
ticket_path['Server']['Name'], # Server
|
||||
size, # Size (duh)
|
||||
ticket_path['Source'].name, # Ticket folder
|
||||
item.name, # Image file
|
||||
),
|
||||
'Server': ticket_path['Server'],
|
||||
'Source': item})
|
||||
# Check for local sources
|
||||
print_standard('Scanning for local sources...')
|
||||
set_thread_error_mode(silent=True) # Prevents "No disk" popups
|
||||
sys_drive = global_vars['Env']['SYSTEMDRIVE']
|
||||
for d in psutil.disk_partitions():
|
||||
if re.search(r'^{}'.format(sys_drive), d.mountpoint, re.IGNORECASE):
|
||||
# Skip current OS drive
|
||||
continue
|
||||
if 'fixed' in d.opts:
|
||||
# Skip DVD, etc
|
||||
sources.append({
|
||||
'Name': '{:9}| File-Based: [DISK] {}'.format(
|
||||
' Local', d.mountpoint),
|
||||
'Source': LocalDisk(d)})
|
||||
set_thread_error_mode(silent=False) # Return to normal
|
||||
|
||||
# Build Menu
|
||||
sources.sort(key=itemgetter('Name'))
|
||||
actions = [{'Name': 'Quit', 'Letter': 'Q'}]
|
||||
|
||||
# Select backup from sources
|
||||
if len(sources) > 0:
|
||||
selection = menu_select(
|
||||
'Which backup are we using?',
|
||||
main_entries=sources,
|
||||
action_entries=actions,
|
||||
disabled_label='DAMAGED')
|
||||
if selection == 'Q':
|
||||
umount_backup_shares()
|
||||
exit_script()
|
||||
else:
|
||||
selected_source = sources[int(selection)-1]['Source']
|
||||
else:
|
||||
print_error('ERROR: No backups found for ticket: {}.'.format(
|
||||
ticket_number))
|
||||
umount_backup_shares()
|
||||
pause("Press Enter to exit...")
|
||||
exit_script()
|
||||
|
||||
# Done
|
||||
return selected_source
|
||||
|
||||
def select_volume(title='Select disk', auto_select=True):
|
||||
"""Select disk from attached disks. returns dict."""
|
||||
actions = [{'Name': 'Quit', 'Letter': 'Q'}]
|
||||
disks = []
|
||||
|
||||
# Build list of disks
|
||||
set_thread_error_mode(silent=True) # Prevents "No disk" popups
|
||||
for d in psutil.disk_partitions():
|
||||
info = {
|
||||
'Disk': d,
|
||||
'Name': d.mountpoint}
|
||||
try:
|
||||
usage = psutil.disk_usage(d.device)
|
||||
free = '{free} / {total} available'.format(
|
||||
free = human_readable_size(usage.free, 2),
|
||||
total = human_readable_size(usage.total, 2))
|
||||
except Exception:
|
||||
# Meh, leaving unsupported destinations out
|
||||
pass
|
||||
# free = 'Unknown'
|
||||
# info['Disabled'] = True
|
||||
else:
|
||||
info['Display Name'] = '{} ({})'.format(info['Name'], free)
|
||||
disks.append(info)
|
||||
set_thread_error_mode(silent=False) # Return to normal
|
||||
|
||||
# Skip menu?
|
||||
if len(disks) == 1 and auto_select:
|
||||
return disks[0]
|
||||
|
||||
# Show menu
|
||||
selection = menu_select(title, main_entries=disks, action_entries=actions)
|
||||
if selection == 'Q':
|
||||
exit_script()
|
||||
else:
|
||||
return disks[int(selection)-1]
|
||||
|
||||
def set_thread_error_mode(silent=True):
|
||||
"""Disable or Enable Windows error message dialogs.
|
||||
|
||||
Disable when scanning for disks to avoid popups for empty cardreaders, etc
|
||||
"""
|
||||
# Code borrowed from: https://stackoverflow.com/a/29075319
|
||||
kernel32 = ctypes.WinDLL('kernel32')
|
||||
|
||||
if silent:
|
||||
kernel32.SetThreadErrorMode(SEM_FAIL, ctypes.byref(SEM_NORMAL))
|
||||
else:
|
||||
kernel32.SetThreadErrorMode(SEM_NORMAL, ctypes.byref(SEM_NORMAL))
|
||||
|
||||
def transfer_source(source_obj, dest_path, selected_items):
|
||||
"""Transfer, or extract, files/folders from source to destination."""
|
||||
if source_obj.is_dir():
|
||||
# Run FastCopy for each selection "group"
|
||||
for group in selected_items:
|
||||
try_and_print(message=group['Message'],
|
||||
function=run_fast_copy, cs='Done',
|
||||
items=group['Items'],
|
||||
dest=group['Destination'])
|
||||
else:
|
||||
if REGEX_WIM_FILE.search(source_obj.name):
|
||||
# Extract files from WIM
|
||||
for group in selected_items:
|
||||
try_and_print(message=group['Message'],
|
||||
function=run_wimextract, cs='Done',
|
||||
source=source_obj.path,
|
||||
items=group['Items'],
|
||||
dest=group['Destination'])
|
||||
else:
|
||||
print_error('ERROR: Unsupported image: {}'.format(source_obj.path))
|
||||
raise GenericError
|
||||
|
||||
def umount_backup_shares():
|
||||
"""Unnount the backup shares regardless of current status."""
|
||||
for server in BACKUP_SERVERS:
|
||||
umount_network_share(server)
|
||||
|
||||
def umount_network_share(server):
|
||||
"""Unnount a network share defined by server."""
|
||||
cmd = r'net use \\{IP}\{Share} /delete'.format(**server)
|
||||
cmd = cmd.split(' ')
|
||||
try:
|
||||
run_program(cmd)
|
||||
except Exception:
|
||||
print_error(r'Failed to umount \\{Name}\{Share}.'.format(**server))
|
||||
sleep(1)
|
||||
else:
|
||||
print_info('Umounted {Name}'.format(**server))
|
||||
server['Mounted'] = False
|
||||
|
||||
def wim_contains(source_path, file_path):
|
||||
"""Check if the WIM contains a file or folder."""
|
||||
_cmd = [
|
||||
global_vars['Tools']['wimlib-imagex'], 'dir',
|
||||
source_path, '1',
|
||||
'--path={}'.format(file_path),
|
||||
'--one-file-only']
|
||||
try:
|
||||
run_program(_cmd)
|
||||
except subprocess.CalledProcessError:
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("This file is not meant to be called directly.")
|
||||
366
Scripts/functions/disk.py
Normal file
366
Scripts/functions/disk.py
Normal file
|
|
@ -0,0 +1,366 @@
|
|||
# Wizard Kit PE: Functions - Disk
|
||||
|
||||
from functions.common import *
|
||||
import partition_uids
|
||||
|
||||
# Regex
|
||||
REGEX_BAD_PARTITION = re.compile(r'(RAW|Unknown)', re.IGNORECASE)
|
||||
REGEX_DISK_GPT = re.compile(
|
||||
r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}',
|
||||
re.IGNORECASE)
|
||||
REGEX_DISK_MBR = re.compile(r'Disk ID: [A-Z0-9]+', re.IGNORECASE)
|
||||
REGEX_DISK_RAW = re.compile(r'Disk ID: 00000000', re.IGNORECASE)
|
||||
|
||||
def assign_volume_letters():
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
for vol in get_volumes():
|
||||
script.write('select volume {}\n'.format(vol['Number']))
|
||||
script.write('assign\n')
|
||||
|
||||
# Remove current letters
|
||||
remove_volume_letters()
|
||||
|
||||
# Run script
|
||||
try:
|
||||
run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
|
||||
def get_boot_mode():
|
||||
boot_mode = 'Legacy'
|
||||
try:
|
||||
reg_key = winreg.OpenKey(
|
||||
winreg.HKEY_LOCAL_MACHINE, r'System\CurrentControlSet\Control')
|
||||
reg_value = winreg.QueryValueEx(reg_key, 'PEFirmwareType')[0]
|
||||
if reg_value == 2:
|
||||
boot_mode = 'UEFI'
|
||||
except:
|
||||
boot_mode = 'Unknown'
|
||||
|
||||
return boot_mode
|
||||
|
||||
def get_disk_details(disk):
|
||||
details = {}
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('detail disk\n')
|
||||
|
||||
try:
|
||||
# Run script
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Remove empty lines
|
||||
tmp = [s.strip() for s in output.splitlines() if s.strip() != '']
|
||||
# Set disk name
|
||||
details['Name'] = tmp[4]
|
||||
# Split each line on ':' skipping those without ':'
|
||||
tmp = [s.split(':') for s in tmp if ':' in s]
|
||||
# Add key/value pairs to the details variable and return dict
|
||||
details.update({key.strip(): value.strip() for (key, value) in tmp})
|
||||
|
||||
return details
|
||||
|
||||
def get_disks():
|
||||
disks = []
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('list disk\n')
|
||||
|
||||
try:
|
||||
# Run script
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Append disk numbers
|
||||
for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', output):
|
||||
num = tmp[0]
|
||||
size = human_readable_size(tmp[1])
|
||||
disks.append({'Number': num, 'Size': size})
|
||||
|
||||
return disks
|
||||
|
||||
def get_partition_details(disk, partition):
|
||||
details = {}
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('select partition {}\n'.format(partition['Number']))
|
||||
script.write('detail partition\n')
|
||||
|
||||
# Diskpart details
|
||||
try:
|
||||
# Run script
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Get volume letter or RAW status
|
||||
tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', output)
|
||||
if tmp:
|
||||
if tmp.group(1).upper() == 'RAW':
|
||||
details['FileSystem'] = RAW
|
||||
else:
|
||||
details['Letter'] = tmp.group(1)
|
||||
# Remove empty lines from output
|
||||
tmp = [s.strip() for s in output.splitlines() if s.strip() != '']
|
||||
# Split each line on ':' skipping those without ':'
|
||||
tmp = [s.split(':') for s in tmp if ':' in s]
|
||||
# Add key/value pairs to the details variable and return dict
|
||||
details.update({key.strip(): value.strip() for (key, value) in tmp})
|
||||
|
||||
# Get MBR type / GPT GUID for extra details on "Unknown" partitions
|
||||
guid = partition_uids.lookup_guid(details['Type'])
|
||||
if guid:
|
||||
details.update({
|
||||
'Description': guid.get('Description', ''),
|
||||
'OS': guid.get('OS', '')})
|
||||
|
||||
if 'Letter' in details:
|
||||
# Disk usage
|
||||
tmp = shutil.disk_usage('{}:\\'.format(details['Letter']))
|
||||
details['Used Space'] = human_readable_size(tmp.used)
|
||||
|
||||
# fsutil details
|
||||
cmd = [
|
||||
'fsutil',
|
||||
'fsinfo',
|
||||
'volumeinfo',
|
||||
'{}:'.format(details['Letter'])
|
||||
]
|
||||
try:
|
||||
output = run_program(cmd)
|
||||
output = output.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Remove empty lines from output
|
||||
tmp = [s.strip() for s in output.splitlines() if s.strip() != '']
|
||||
# Add "Feature" lines
|
||||
details['File System Features'] = [s.strip() for s in tmp
|
||||
if ':' not in s]
|
||||
# Split each line on ':' skipping those without ':'
|
||||
tmp = [s.split(':') for s in tmp if ':' in s]
|
||||
# Add key/value pairs to the details variable and return dict
|
||||
details.update({key.strip(): value.strip() for (key, value) in tmp})
|
||||
|
||||
# Set Volume Name
|
||||
details['Name'] = details.get('Volume Name', '')
|
||||
|
||||
# Set FileSystem Type
|
||||
if details.get('FileSystem', '') != 'RAW':
|
||||
details['FileSystem'] = details.get('File System Name', 'Unknown')
|
||||
|
||||
return details
|
||||
|
||||
def get_partitions(disk):
|
||||
partitions = []
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('list partition\n')
|
||||
|
||||
try:
|
||||
# Run script
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Append partition numbers
|
||||
regex = r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+'
|
||||
for tmp in re.findall(regex, output, re.IGNORECASE):
|
||||
num = tmp[0]
|
||||
size = human_readable_size(tmp[1])
|
||||
partitions.append({'Number': num, 'Size': size})
|
||||
|
||||
return partitions
|
||||
|
||||
def get_table_type(disk):
|
||||
part_type = 'Unknown'
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('uniqueid disk\n')
|
||||
|
||||
try:
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
if REGEX_DISK_GPT.search(output):
|
||||
part_type = 'GPT'
|
||||
elif REGEX_DISK_MBR.search(output):
|
||||
part_type = 'MBR'
|
||||
elif REGEX_DISK_RAW.search(output):
|
||||
part_type = 'RAW'
|
||||
else:
|
||||
part_type = 'Unknown
|
||||
|
||||
return part_type
|
||||
|
||||
def get_volumes():
|
||||
vols = []
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('list volume\n')
|
||||
|
||||
try:
|
||||
# Run script
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Append volume numbers
|
||||
for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', output):
|
||||
vols.append({'Number': tmp[0], 'Letter': tmp[1]})
|
||||
|
||||
return vols
|
||||
|
||||
def is_bad_partition(par):
|
||||
return 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem'])
|
||||
|
||||
def prep_disk_for_formatting(disk=None):
|
||||
disk['Format Warnings'] = '\n'
|
||||
width = len(str(len(disk['Partitions'])))
|
||||
|
||||
# Bail early
|
||||
if disk is None:
|
||||
raise Exception('Disk not provided.')
|
||||
|
||||
# Set boot method and partition table type
|
||||
disk['Use GPT'] = True
|
||||
if (get_boot_mode() == 'UEFI'):
|
||||
if (not ask("Setup Windows to use UEFI booting?")):
|
||||
disk['Use GPT'] = False
|
||||
else:
|
||||
if (ask("Setup Windows to use BIOS/Legacy booting?")):
|
||||
disk['Use GPT'] = False
|
||||
|
||||
# Set Display and Warning Strings
|
||||
if len(disk['Partitions']) == 0:
|
||||
disk['Format Warnings'] += 'No partitions found\n'
|
||||
for partition in disk['Partitions']:
|
||||
display = ' Partition {num:>{width}}:\t{size} {fs}'.format(
|
||||
num = partition['Number'],
|
||||
width = width,
|
||||
size = partition['Size'],
|
||||
fs = partition['FileSystem'])
|
||||
|
||||
if is_bad_partition(partition):
|
||||
# Set display string using partition description & OS type
|
||||
display += '\t\t{q}{name}{q}\t{desc} ({os})'.format(
|
||||
display = display,
|
||||
q = '"' if partition['Name'] != '' else '',
|
||||
name = partition['Name'],
|
||||
desc = partition['Description'],
|
||||
os = partition['OS'])
|
||||
else:
|
||||
# List space used instead of partition description & OS type
|
||||
display += ' (Used: {used})\t{q}{name}{q}'.format(
|
||||
used = partition['Used Space'],
|
||||
q = '"' if partition['Name'] != '' else '',
|
||||
name = partition['Name'])
|
||||
# For all partitions
|
||||
partition['Display String'] = display
|
||||
|
||||
def reassign_volume_letter(letter, new_letter='I'):
|
||||
if not letter:
|
||||
# Ignore
|
||||
return None
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select volume {}\n'.format(letter))
|
||||
script.write('remove noerr\n')
|
||||
script.write('assign letter={}\n'.format(new_letter))
|
||||
run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
return new_letter
|
||||
|
||||
def remove_volume_letters(keep=None):
|
||||
if not keep:
|
||||
keep = ''
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
for vol in get_volumes():
|
||||
if vol['Letter'].upper() != keep.upper():
|
||||
script.write('select volume {}\n'.format(vol['Number']))
|
||||
script.write('remove noerr\n')
|
||||
|
||||
# Run script
|
||||
try:
|
||||
run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
|
||||
def scan_disks():
|
||||
"""Get details about the attached disks"""
|
||||
disks = get_disks()
|
||||
|
||||
# Get disk details
|
||||
for disk in disks:
|
||||
# Get partition style
|
||||
disk['Table'] = get_table_type(disk)
|
||||
|
||||
# Get disk name/model and physical details
|
||||
disk.update(get_disk_details(disk))
|
||||
|
||||
# Get partition info for disk
|
||||
disk['Partitions'] = get_partitions(disk)
|
||||
|
||||
for partition in disk['Partitions']:
|
||||
# Get partition details
|
||||
partition.update(get_partition_details(disk, partition))
|
||||
|
||||
# Done
|
||||
return disks
|
||||
|
||||
def select_disk(title='Which disk?', disks):
|
||||
"""Select a disk from the attached disks"""
|
||||
# Build menu
|
||||
disk_options = []
|
||||
for disk in disks:
|
||||
display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk)
|
||||
pwidth=len(str(len(disk['Partitions'])))
|
||||
for partition in disk['Partitions']:
|
||||
# Main text
|
||||
p_name = 'Partition {num:>{width}}: {size} ({fs})'.format(
|
||||
num = partition['Number'],
|
||||
width = pwidth,
|
||||
size = partition['Size'],
|
||||
fs = partition['FileSystem'])
|
||||
if partition['Name']:
|
||||
p_name += '\t"{}"'.format(partition['Name'])
|
||||
|
||||
# Show unsupported partition(s)
|
||||
if is_bad_partition(partition):
|
||||
p_display_name = '{YELLOW}{display}{CLEAR}'.format(
|
||||
display=p_name, **COLORS)
|
||||
|
||||
display_name += '\n\t\t\t{}'.format(display_name)
|
||||
if not disk['Partitions']:
|
||||
display_name += '\n\t\t\t{}No partitions found.{}'.format(
|
||||
COLORS['YELLOW'], COLORS['CLEAR'])
|
||||
|
||||
disk_options.append({'Name': display_name, 'Disk': disk})
|
||||
actions = [
|
||||
{'Name': 'Main Menu', 'Letter': 'M'},
|
||||
]
|
||||
|
||||
# Menu loop
|
||||
selection = menu_select(
|
||||
title = title,
|
||||
main_entries = disk_options,
|
||||
action_entries = actions)
|
||||
|
||||
if (selection.isnumeric()):
|
||||
return disk_options[int(selection)-1]['Disk']
|
||||
elif (selection == 'M'):
|
||||
raise GeneralAbort
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("This file is not meant to be called directly.")
|
||||
|
|
@ -1,4 +1,4 @@
|
|||
# WK WinPE - PARTITION UIDs
|
||||
# Wizard Kit PE: Functions - PARTITION UIDs
|
||||
# sources: https://en.wikipedia.org/wiki/GUID_Partition_Table
|
||||
# https://en.wikipedia.org/wiki/Partition_type
|
||||
|
||||
225
Scripts/functions/windows_setup.py
Normal file
225
Scripts/functions/windows_setup.py
Normal file
|
|
@ -0,0 +1,225 @@
|
|||
# Wizard Kit PE: Functions - Windows Setup
|
||||
|
||||
from functions.data import *
|
||||
|
||||
# STATIC VARIABLES
|
||||
DISKPART_SCRIPT = r'{}\diskpart.script'.format(global_vars['Env']['TMP'])
|
||||
WINDOWS_VERSIONS = [
|
||||
{'Name': 'Windows 7 Home Basic',
|
||||
'Image File': 'Win7',
|
||||
'Image Name': 'Windows 7 HOMEBASIC'},
|
||||
{'Name': 'Windows 7 Home Premium',
|
||||
'Image File': 'Win7',
|
||||
'Image Name': 'Windows 7 HOMEPREMIUM'},
|
||||
{'Name': 'Windows 7 Professional',
|
||||
'Image File': 'Win7',
|
||||
'Image Name': 'Windows 7 PROFESSIONAL'},
|
||||
{'Name': 'Windows 7 Ultimate',
|
||||
'Image File': 'Win7',
|
||||
'Image Name': 'Windows 7 ULTIMATE'},
|
||||
|
||||
{'Name': 'Windows 8.1',
|
||||
'Image File': 'Win8',
|
||||
'Image Name': 'Windows 8.1',
|
||||
'CRLF': True},
|
||||
{'Name': 'Windows 8.1 Pro',
|
||||
'Image File': 'Win8',
|
||||
'Image Name': 'Windows 8.1 Pro'},
|
||||
|
||||
{'Name': 'Windows 10 Home',
|
||||
'Image File': 'Win10',
|
||||
'Image Name': 'Windows 10 Home',
|
||||
'CRLF': True},
|
||||
{'Name': 'Windows 10 Pro',
|
||||
'Image File': 'Win10',
|
||||
'Image Name': 'Windows 10 Pro'},
|
||||
]
|
||||
|
||||
def find_windows_image(windows_version):
|
||||
"""Search for a Windows source image file, returns dict.
|
||||
|
||||
Searches on local disks and then the WINDOWS_SERVER share."""
|
||||
image = {}
|
||||
imagefile = windows_version['Image File']
|
||||
imagename = windows_version['Image Name']
|
||||
|
||||
# Search local source
|
||||
for d in psutil.disk_partitions():
|
||||
for ext in ['esd', 'wim', 'swm']:
|
||||
path = '{}images\{}.{}'.format(d.mountpoint, imagefile, ext)
|
||||
if os.path.isfile(path) and wim_contains_image(path, imagename):
|
||||
image['Path'] = path
|
||||
image['Source'] = letter
|
||||
if ext == 'swm':
|
||||
image['Glob'] = '--ref="{}*.swm"'.format(image['Path'][:-4])
|
||||
break
|
||||
|
||||
# Check for network source
|
||||
if not image:
|
||||
mount_windows_share()
|
||||
if not WINDOWS_SERVER['Mounted']:
|
||||
return None
|
||||
for ext in ['esd', 'wim', 'swm']:
|
||||
path = r'\\{}\{}\images\{}.ext'.format(
|
||||
WINDOWS_SERVER['IP'], WINDOWS_SERVER['Share'], imagefile, ext)
|
||||
if os.path.isfile(path) and wim_contains_image(path, imagename):
|
||||
image['Path'] = path
|
||||
image['Source'] = None
|
||||
if ext == 'swm':
|
||||
image['Glob'] = '--ref="{}*.swm"'.format(image['Path'][:-4])
|
||||
break
|
||||
|
||||
# Display image to be used (if any) and return
|
||||
if image:
|
||||
print_info('Using image: {}'.format(image['Path']))
|
||||
return image
|
||||
else:
|
||||
print_error('Failed to find Windows source image for {}'.format(
|
||||
windows_version['Name']))
|
||||
raise GeneralAbort
|
||||
|
||||
def format_disk(disk, use_gpt):
|
||||
"""Format disk for use as a Windows OS disk."""
|
||||
if use_gpt:
|
||||
format_gpt(disk)
|
||||
else:
|
||||
format_mbr(disk)
|
||||
|
||||
def format_gpt(disk):
|
||||
"""Format disk for use as a Windows OS disk using the GPT layout."""
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
# Partition table
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('clean\n')
|
||||
script.write('convert gpt\n')
|
||||
|
||||
# System partition
|
||||
# NOTE: ESP needs to be >= 260 for Advanced Format 4K disks
|
||||
script.write('create partition efi size=500\n')
|
||||
script.write('format quick fs=fat32 label="System"\n')
|
||||
script.write('assign letter="S"\n')
|
||||
|
||||
# Microsoft Reserved (MSR) partition
|
||||
script.write('create partition msr size=128\n')
|
||||
|
||||
# Windows partition
|
||||
script.write('create partition primary\n')
|
||||
script.write('format quick fs=ntfs label="Windows"\n')
|
||||
script.write('assign letter="W"\n')
|
||||
|
||||
# Recovery Tools partition
|
||||
script.write('shrink minimum=500\n')
|
||||
script.write('create partition primary\n')
|
||||
script.write('format quick fs=ntfs label="Recovery Tools"\n')
|
||||
script.write('assign letter="T"\n')
|
||||
script.write('set id="de94bba4-06d1-4d40-a16a-bfd50179d6ac"\n')
|
||||
script.write('gpt attributes=0x8000000000000001\n')
|
||||
|
||||
# Run script
|
||||
run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
time.sleep(2)
|
||||
|
||||
def format_mbr(disk):
|
||||
"""Format disk for use as a Windows OS disk using the MBR layout."""
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
# Partition table
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('clean\n')
|
||||
|
||||
# System partition
|
||||
script.write('create partition primary size=100\n')
|
||||
script.write('format fs=ntfs quick label="System Reserved"\n')
|
||||
script.write('active\n')
|
||||
script.write('assign letter="S"\n')
|
||||
|
||||
# Windows partition
|
||||
script.write('create partition primary\n')
|
||||
script.write('format fs=ntfs quick label="Windows"\n')
|
||||
script.write('assign letter="W"\n')
|
||||
|
||||
# Recovery Tools partition
|
||||
script.write('shrink minimum=500\n')
|
||||
script.write('create partition primary\n')
|
||||
script.write('format quick fs=ntfs label="Recovery"\n')
|
||||
script.write('assign letter="T"\n')
|
||||
script.write('set id=27\n')
|
||||
|
||||
# Run script
|
||||
run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
time.sleep(2)
|
||||
|
||||
def mount_windows_share():
|
||||
"""Mount the Windows images share unless labeled as already mounted."""
|
||||
if WINDOWS_SERVER['Mounted']:
|
||||
# Blindly skip if we mounted earlier
|
||||
continue
|
||||
|
||||
mount_network_share(WINDOWS_SERVER)
|
||||
|
||||
def select_windows_version():
|
||||
actions = [
|
||||
{'Name': 'Main Menu', 'Letter': 'M'},
|
||||
]
|
||||
|
||||
# Menu loop
|
||||
selection = menu_select(
|
||||
title = 'Which version of Windows are we installing?',
|
||||
main_entries = WINDOWS_VERSIONS,
|
||||
action_entries = actions)
|
||||
|
||||
if selection.isnumeric():
|
||||
return WINDOWS_VERSIONS[int(selection)-1]
|
||||
elif selection == 'M':
|
||||
raise GeneralAbort
|
||||
|
||||
def setup_windows(windows_image, windows_version):
|
||||
cmd = [
|
||||
global_vars['Tools']['wimlib-imagex'],
|
||||
'apply',
|
||||
windows_image['Path'],
|
||||
windows_version['Image Name'],
|
||||
'W:\\']
|
||||
if 'Glob' in windows_image:
|
||||
cmd.extend(windows_image['Glob'])
|
||||
run_program(cmd)
|
||||
|
||||
def setup_windows_re(windows_version, windows_letter='W', tools_letter='T'):
|
||||
win = r'{}:\Windows'.format(windows_letter)
|
||||
winre = r'{}\System32\Recovery\WinRE.wim'.format(win)
|
||||
dest = r'{}:\Recovery\WindowsRE'.format(tools_letter)
|
||||
|
||||
# Copy WinRE.wim
|
||||
os.makedirs(dest, exist_ok=True)
|
||||
shutil.copy(winre, r'{}\WinRE.wim'.format(dest))
|
||||
|
||||
# Set location
|
||||
cmd = [
|
||||
r'{}\System32\ReAgentc.exe'.format(win),
|
||||
'/setreimage',
|
||||
'/path', dest,
|
||||
'/target', win]
|
||||
run_program(cmd)
|
||||
|
||||
def update_boot_partition(system_letter='S', windows_letter='W', mode='ALL'):
|
||||
cmd = [
|
||||
r'{}:\Windows\System32\bcdboot.exe'.format(windows_letter),
|
||||
r'{}:\Windows'.format(windows_letter),
|
||||
'/s', '{}:'.format(system_letter),
|
||||
'/f', mode]
|
||||
run_program(cmd)
|
||||
|
||||
def wim_contains_image(filename, imagename):
|
||||
cmd = [
|
||||
global_vars['Tools']['wimlib-imagex'],
|
||||
'info',
|
||||
filename,
|
||||
imagename]
|
||||
try:
|
||||
run_program(cmd)
|
||||
except subprocess.CalledProcessError:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("This file is not meant to be called directly.")
|
||||
380
Scripts/functions/winpe_menus.py
Normal file
380
Scripts/functions/winpe_menus.py
Normal file
|
|
@ -0,0 +1,380 @@
|
|||
# Wizard Kit PE: Menus
|
||||
|
||||
from functions.backup import *
|
||||
from functions.disk import *
|
||||
from functions.windows_setup import *
|
||||
|
||||
# STATIC VARIABLES
|
||||
FAST_COPY_PE_ARGS = [
|
||||
'/cmd=noexist_only',
|
||||
'/utf8',
|
||||
'/skip_empty_dir',
|
||||
'/linkdest',
|
||||
'/no_ui',
|
||||
'/auto_close',
|
||||
'/exclude={}'.format(';'.join(FAST_COPY_EXCLUDES)),
|
||||
]
|
||||
PE_TOOLS = {
|
||||
'BlueScreenView': {
|
||||
'Path': r'BlueScreenView\BlueScreenView.exe',
|
||||
},
|
||||
'FastCopy': {
|
||||
'Path': r'FastCopy\FastCopy.exe',
|
||||
'Args': FAST_COPY_PE_ARGS,
|
||||
},
|
||||
'HWiNFO': {
|
||||
'Path': r'HWiNFO\HWiNFO.exe',
|
||||
},
|
||||
'NT Password Editor': {
|
||||
'Path': r'NT Password Editor\ntpwedit.exe',
|
||||
},
|
||||
'Notepad++': {
|
||||
'Path': r'NotepadPlusPlus\NotepadPlusPlus.exe',
|
||||
},
|
||||
'PhotoRec': {
|
||||
'Path': r'TestDisk\photorec_win.exe',
|
||||
'Args': ['-new_console:n'],
|
||||
},
|
||||
'Prime95': {
|
||||
'Path': r'Prime95\prime95.exe',
|
||||
},
|
||||
'ProduKey': {
|
||||
'Path': r'ProduKey\ProduKey.exe',
|
||||
},
|
||||
'Q-Dir': {
|
||||
'Path': r'Q-Dir\Q-Dir.exe',
|
||||
},
|
||||
'TestDisk': {
|
||||
'Path': r'TestDisk\testdisk_win.exe',
|
||||
'Args': ['-new_console:n'],
|
||||
},
|
||||
}
|
||||
|
||||
def menu_backup():
|
||||
"""Take backup images of partition(s) in the WIM format."""
|
||||
errors = False
|
||||
other_results = {
|
||||
'Error': {
|
||||
'CalledProcessError': 'Unknown Error',
|
||||
'PathNotFoundError': 'Missing',
|
||||
},
|
||||
'Warning': {
|
||||
'GenericAbort': 'Skipped',
|
||||
'GenericRepair': 'Repaired',
|
||||
}}
|
||||
set_title('{}: Backup Menu'.format(KIT_NAME_FULL))
|
||||
|
||||
# Set ticket Number
|
||||
clear_screen()
|
||||
ticket_number = get_ticket_number()
|
||||
|
||||
# Mount backup shares
|
||||
mount_backup_shares()
|
||||
|
||||
# Select destination
|
||||
destination = select_backup_destination()
|
||||
|
||||
# Scan disks
|
||||
try_and_print(
|
||||
message = 'Assigning letters...',
|
||||
function = assign_volume_letters,
|
||||
other_results = other_results)
|
||||
result = try_and_print(
|
||||
message = 'Getting disk info...',
|
||||
function = scan_disks,
|
||||
other_results = other_results)
|
||||
if result['CS']:
|
||||
disks = result['Out']
|
||||
else:
|
||||
print_error('ERROR: No disks found.')
|
||||
raise GenericAbort
|
||||
|
||||
# Select disk to backup
|
||||
disk = select_disk('For which disk are we creating backups?', disks)
|
||||
if not disk:
|
||||
raise GenericAbort
|
||||
|
||||
# "Prep" disk
|
||||
prep_disk_for_backup(destination, disk, ticket_number)
|
||||
|
||||
# Display details for backup task
|
||||
clear_screen()
|
||||
print_info('Create Backup - Details:\n')
|
||||
show_info(message='Ticket:', info=ticket_number)
|
||||
show_info(
|
||||
message = 'Source:',
|
||||
info = '[{Table}] ({Type}) {Name} {Size}'.format(**disk),
|
||||
)
|
||||
show_info(
|
||||
message = 'Destination:',
|
||||
info = destination.get('Display Name', destination['Name']),
|
||||
)
|
||||
for par in disk['Partitions']:
|
||||
show_info(message='', info=par['Display String'], width=20)
|
||||
print_standard(disk['Backup Warnings'])
|
||||
|
||||
# Ask to proceed
|
||||
if (not ask('Proceed with backup?')):
|
||||
raise GenericAbort
|
||||
|
||||
# Backup partition(s)
|
||||
print_info('\n\nStarting task.\n')
|
||||
for par in disk['Partitions']:
|
||||
result = try_and_print(
|
||||
message = 'Partition {} Backup...'.format(par['Number']),
|
||||
function = backup_partition,
|
||||
other_results = other_results,
|
||||
disk = disk,
|
||||
partition = par)
|
||||
if not result['CS']:
|
||||
errors = True
|
||||
par['Error'] = result['Error']
|
||||
|
||||
# Verify backup(s)
|
||||
if disk['Valid Partitions']:
|
||||
print_info('\n\n Verifying backup images(s)\n')
|
||||
for par in disk['Partitions']:
|
||||
if par['Number'] in disk['Bad Partitions']:
|
||||
continue # Skip verification
|
||||
result = try_and_print(
|
||||
message = 'Partition {} Image...'.format(par['Number']),
|
||||
function = verify_wim_backup,
|
||||
other_results = other_results,
|
||||
partition = par)
|
||||
if not result['CS']:
|
||||
errors = True
|
||||
par['Error'] = result['Error']
|
||||
|
||||
# Print summary
|
||||
if errors:
|
||||
print_warning('\nErrors were encountered and are detailed below.')
|
||||
for par in [p for p in disk['Partitions'] if 'Error' in p]:
|
||||
print_standard(' Partition {} Error:'.format(par['Number']))
|
||||
if hasattr(par['Error'], 'stderr'):
|
||||
try:
|
||||
par['Error'] = par['Error'].stderr.decode()
|
||||
except:
|
||||
# Deal with badly formatted error message
|
||||
pass
|
||||
if isinstance(par['Error'], basestring):
|
||||
print_error('\t{}'.format(par['Error']))
|
||||
else:
|
||||
try:
|
||||
par['Error'] = par['Error'].splitlines()
|
||||
par['Error'] = [line.strip() for line in par['Error']]
|
||||
par['Error'] = [line for line in par['Error'] if line]
|
||||
except:
|
||||
pass
|
||||
for line in par['Error']:
|
||||
print_error('\t{}'.format(line))
|
||||
time.sleep(30)
|
||||
else:
|
||||
print_success('\nNo errors were encountered during imaging.')
|
||||
time.sleep(5)
|
||||
pause('\nPress Enter to return to main menu... ')
|
||||
|
||||
def menu_root():
|
||||
menus = [
|
||||
{'Name': 'Create Backups', 'Menu': menu_backup},
|
||||
{'Name': 'Setup Windows', 'Menu': menu_setup},
|
||||
{'Name': 'Misc Tools', 'Menu': menu_tools},
|
||||
]
|
||||
actions = [
|
||||
{'Name': 'Command Prompt', 'Letter': 'C'},
|
||||
{'Name': 'Reboot', 'Letter': 'R'},
|
||||
{'Name': 'Shutdown', 'Letter': 'S'},
|
||||
]
|
||||
|
||||
# Main loop
|
||||
while True:
|
||||
set_title(KIT_NAME_FULL)
|
||||
selection = menu_select(
|
||||
title = 'Main Menu',
|
||||
main_entries = menus,
|
||||
action_entries = actions,
|
||||
secret_exit = True)
|
||||
|
||||
if (selection.isnumeric()):
|
||||
try:
|
||||
menus[int(selection)-1]['Menu']()
|
||||
except AbortError:
|
||||
pass
|
||||
elif (selection == 'C'):
|
||||
run_program(['cmd', '-new_console:n'], check=False)
|
||||
elif (selection == 'R'):
|
||||
run_program(['wpeutil', 'reboot'])
|
||||
elif (selection == 'S'):
|
||||
run_program(['wpeutil', 'shutdown'])
|
||||
else:
|
||||
exit_script()
|
||||
|
||||
def menu_setup():
|
||||
"""Format a disk (MBR/GPT), apply a Windows image, and setup boot files."""
|
||||
errors = False
|
||||
set_title('{}: Setup Menu'.format(KIT_NAME_FULL))
|
||||
|
||||
# Set ticket ID
|
||||
clear_screen()
|
||||
ticket_number = get_ticket_number()
|
||||
|
||||
# Select the version of Windows to apply
|
||||
windows_version = select_windows_version()
|
||||
|
||||
# Find Windows image
|
||||
windows_image = find_windows_image(windows_version)
|
||||
|
||||
# Scan disks
|
||||
try_and_print(
|
||||
message = 'Assigning letters...',
|
||||
function = assign_volume_letters,
|
||||
other_results = other_results)
|
||||
result = try_and_print(
|
||||
message = 'Getting disk info...',
|
||||
function = scan_disks,
|
||||
other_results = other_results)
|
||||
if result['CS']:
|
||||
disks = result['Out']
|
||||
else:
|
||||
print_error('ERROR: No disks found.')
|
||||
raise GenericAbort
|
||||
|
||||
# Select disk to use as the OS disk
|
||||
dest_disk = select_disk('To which disk are we installing Windows?', disks)
|
||||
if not disk:
|
||||
raise GenericAbort
|
||||
|
||||
# "Prep" disk
|
||||
prep_disk_for_formatting(dest_disk)
|
||||
|
||||
# Display details for setup task
|
||||
clear_screen()
|
||||
print_info('Setup Windows - Details:\n')
|
||||
show_info(message='Ticket:', info=ticket_number)
|
||||
show_info(message='Installing:', info=windows_version['Name'])
|
||||
show_info(
|
||||
message = 'Boot Method:',
|
||||
info = 'UEFI (GPT)' if dest_disk['Use GPT'] else 'Legacy (MBR)')
|
||||
show_info(message='Using Image:', info=windows_version['Path'])
|
||||
print_warning(' ERASING: \t[{Table}] ({Type}) {Name} {Size}\n'.format(
|
||||
**dest_disk))
|
||||
for par in dest_disk['Partitions']:
|
||||
print_warning(par['Display String'])
|
||||
print_warning(dest_disk['Format Warnings'])
|
||||
|
||||
if (not ask('Is this correct?')):
|
||||
raise GeneralAbort
|
||||
|
||||
# Safety check
|
||||
print_standard('\nSAFETY CHECK')
|
||||
print_warning('All data will be DELETED from the '
|
||||
'disk & partition(s) listed above.')
|
||||
print_warning('This is irreversible and will lead '
|
||||
'to {CLEAR}{RED}DATA LOSS.'.format(**COLORS))
|
||||
if (not ask('Asking again to confirm, is this correct?')):
|
||||
raise GeneralAbort
|
||||
|
||||
# Remove volume letters so S, T, & W can be used below
|
||||
remove_volume_letters(keep=windows_image['Source'])
|
||||
new_letter = reassign_volume_letter(letter=windows_image['Source'])
|
||||
if new_letter:
|
||||
windows_image['Source'] = new_letter
|
||||
|
||||
# Format and partition disk
|
||||
result = try_and_print(
|
||||
message = 'Formatting disk...',
|
||||
function = format_disk,
|
||||
other_results = other_results,
|
||||
disk = dest_disk,
|
||||
use_gpt = dest_disk['Use GPT'])
|
||||
if not result['CS']:
|
||||
# We need to crash as the disk is in an unknown state
|
||||
print_error('ERROR: Failed to format disk.')
|
||||
raise GenericAbort
|
||||
|
||||
# Apply Image
|
||||
result = try_and_print(
|
||||
message = 'Applying image...',
|
||||
function = setup_windows,
|
||||
other_results = other_results,
|
||||
windows_image = windows_image,
|
||||
windows_version = windows_version)
|
||||
if not result['CS']:
|
||||
# We need to crash as the disk is in an unknown state
|
||||
print_error('ERROR: Failed to apply image.')
|
||||
raise GenericAbort
|
||||
|
||||
# Create Boot files
|
||||
try_and_print(
|
||||
message = 'Updating boot files...',
|
||||
function = update_boot_partition,
|
||||
other_results = other_results)
|
||||
|
||||
# Setup WinRE
|
||||
try_and_print(
|
||||
message = 'Updating recovery tools...',
|
||||
function = setup_windows_re,
|
||||
other_results = other_results,
|
||||
windows_version = windows_version)
|
||||
|
||||
# Print summary
|
||||
print_standard('\nDone.')
|
||||
pause('\nPress Enter to return to main menu... ')
|
||||
|
||||
def menu_tools():
|
||||
tools = [k for k in sorted(PE_TOOLS.keys())]
|
||||
actions = [{'Name': 'Main Menu', 'Letter': 'M'},]
|
||||
set_title(KIT_NAME_FULL)
|
||||
|
||||
# Menu loop
|
||||
while True:
|
||||
selection = menu_select(
|
||||
title = 'Tools Menu',
|
||||
main_entries = tools,
|
||||
action_entries = actions)
|
||||
if (selection.isnumeric()):
|
||||
tool = tools[int(selection)-1]
|
||||
cmd = [PE_TOOLS[tool]['Path']] + PE_TOOLS[tool].get('Args', [])
|
||||
if tool == 'Blue Screen View':
|
||||
# Select path to scan
|
||||
minidump_path = select_minidump_path()
|
||||
if minidump_path:
|
||||
cmd.extend(['/MiniDumpFolder', minidump_path])
|
||||
try:
|
||||
popen_program(cmd)
|
||||
except Exception:
|
||||
print_error('Failed to run {prog}'.format(prog=tool['Name']))
|
||||
time.sleep(2)
|
||||
pause()
|
||||
elif (selection == 'M'):
|
||||
break
|
||||
|
||||
def select_minidump_path():
|
||||
dumps = []
|
||||
|
||||
# Assign volume letters first
|
||||
assign_volume_letters()
|
||||
|
||||
# Search for minidumps
|
||||
tmp = run_program('mountvol')
|
||||
tmp = [d for d in re.findall(r'.*([A-Za-z]):\\', tmp.stdout.decode())]
|
||||
# Remove RAMDisk letter
|
||||
if 'X' in tmp:
|
||||
tmp.remove('X')
|
||||
for disk in tmp:
|
||||
if os.path.exists('{}:\\Windows\\MiniDump'.format(disk)):
|
||||
dumps.append({'Name': '{}:\\Windows\\MiniDump'.format(disk)})
|
||||
|
||||
# Check results before showing menu
|
||||
if len(dumps) == 0:
|
||||
print_error(' No BSoD / MiniDump paths found')
|
||||
time.sleep(2)
|
||||
return None
|
||||
|
||||
# Menu
|
||||
selection = menu_select(
|
||||
title = 'Which BSoD / MiniDump path are we scanning?',
|
||||
main_entries = dumps)
|
||||
return dumps[int(selection) - 1]['Name']
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("This file is not meant to be called directly.")
|
||||
276
Scripts/menu.py
276
Scripts/menu.py
|
|
@ -1,276 +0,0 @@
|
|||
# WK WinPE Menu
|
||||
|
||||
# Init
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
os.chdir(os.path.dirname(os.path.realpath(__file__)))
|
||||
bin = os.path.abspath('..\\')
|
||||
sys.path.append(os.getcwd())
|
||||
from functions import *
|
||||
|
||||
## Colors
|
||||
COLORS = {
|
||||
'CLEAR': '\033[0m',
|
||||
'RED': '\033[31m',
|
||||
'GREEN': '\033[32m',
|
||||
'YELLOW': '\033[33m',
|
||||
'BLUE': '\033[34m'}
|
||||
|
||||
def menu_backup_imaging():
|
||||
"""Take backup images of partition(s) in the WIM format and save them to a backup share"""
|
||||
errors = False
|
||||
|
||||
# Set ticket ID
|
||||
os.system('cls')
|
||||
ticket_id = get_ticket_id()
|
||||
|
||||
# Mount backup shares
|
||||
mount_backup_shares()
|
||||
|
||||
# Select destination
|
||||
dest = select_destination()
|
||||
if dest is None:
|
||||
abort_to_main_menu('Aborting Backup Creation')
|
||||
|
||||
# Select disk to backup
|
||||
disk = select_disk('For which drive are we creating backups?')
|
||||
if disk is None:
|
||||
abort_to_main_menu()
|
||||
prep_disk_for_backup(dest, disk, ticket_id)
|
||||
|
||||
# Display details for backup task
|
||||
os.system('cls')
|
||||
print('Create Backup - Details:\n')
|
||||
print(' Ticket: \t{ticket_id}'.format(ticket_id=ticket_id))
|
||||
print(' Source: \t[{Table}] ({Type}) {Name} {Size}\n'.format(**disk))
|
||||
print(' Destination:\t{name}'.format(name=dest.get('Display Name', dest['Name'])))
|
||||
for par in disk['Partitions']:
|
||||
print(par['Display String'])
|
||||
print(disk['Backup Warnings'])
|
||||
|
||||
# Ask to proceed
|
||||
if (not ask('Proceed with backup?')):
|
||||
abort_to_main_menu('Aborting Backup Creation')
|
||||
|
||||
# Backup partition(s)
|
||||
print('\n\nStarting task.\n')
|
||||
for par in disk['Partitions']:
|
||||
try:
|
||||
backup_partition(bin, disk, par)
|
||||
except BackupError:
|
||||
errors = True
|
||||
|
||||
# Verify backup(s)
|
||||
if disk['Valid Partitions'] > 1:
|
||||
print('\n\n Verifying backups\n')
|
||||
else:
|
||||
print('\n\n Verifying backup\n')
|
||||
for par in disk['Partitions']:
|
||||
if par['Number'] in disk['Bad Partitions']:
|
||||
continue # Skip verification
|
||||
try:
|
||||
verify_wim_backup(bin, par)
|
||||
except BackupError:
|
||||
errors = True
|
||||
|
||||
# Print summary
|
||||
if errors:
|
||||
print_warning('\nErrors were encountered and are detailed below.')
|
||||
for par in [p for p in disk['Partitions'] if 'Error' in p]:
|
||||
print(' Partition {Number} Error:'.format(**par))
|
||||
for line in [line.strip() for line in par['Error'] if line.strip() != '']:
|
||||
print_error('\t{line}'.format(line=line))
|
||||
time.sleep(30)
|
||||
else:
|
||||
print_success('\nNo errors were encountered during imaging.')
|
||||
time.sleep(5)
|
||||
pause('\nPress Enter to return to main menu... ')
|
||||
|
||||
def menu_windows_setup():
|
||||
"""Format a drive, partition for MBR or GPT, apply a Windows image, and rebuild the boot files"""
|
||||
errors = False
|
||||
|
||||
# Set ticket ID
|
||||
os.system('cls')
|
||||
ticket_id = get_ticket_id()
|
||||
|
||||
# Select the version of Windows to apply
|
||||
windows_version = select_windows_version()
|
||||
|
||||
# Select drive to use as the OS drive
|
||||
dest_disk = select_disk('To which drive are we installing Windows?')
|
||||
prep_disk_for_formatting(dest_disk)
|
||||
|
||||
# Find Windows image
|
||||
## NOTE: Needs to happen AFTER select_disk() is called as there's a hidden assign_volume_letters().
|
||||
## This changes the current letters thus preventing installing from a local source.
|
||||
windows_image = find_windows_image(bin, windows_version)
|
||||
|
||||
# Display details for setup task
|
||||
os.system('cls')
|
||||
print('Setup Windows - Details:\n')
|
||||
print(' Ticket: \t{ticket_id}'.format(ticket_id=ticket_id))
|
||||
print(' Installing: \t{winver}'.format(winver=windows_version['Name']))
|
||||
print(' Boot Method:\t{_type}'.format(
|
||||
_type='UEFI (GPT)' if dest_disk['Use GPT'] else 'Legacy (MBR)'))
|
||||
print(' Using Image:\t{File}.{Ext}'.format(**windows_image))
|
||||
print_warning(' ERASING: \t[{Table}] ({Type}) {Name} {Size}\n'.format(**dest_disk))
|
||||
for par in dest_disk['Partitions']:
|
||||
print_warning(par['Display String'])
|
||||
print_warning(dest_disk['Format Warnings'])
|
||||
|
||||
if (not ask('Is this correct?')):
|
||||
abort_to_main_menu('Aborting Windows setup')
|
||||
|
||||
# Safety check
|
||||
print('\nSAFETY CHECK')
|
||||
print_warning('All data will be DELETED from the drive and partition(s) listed above.')
|
||||
print_warning('This is irreversible and will lead to {CLEAR}{RED}DATA LOSS.'.format(**COLORS))
|
||||
if (not ask('Asking again to confirm, is this correct?')):
|
||||
abort_to_main_menu('Aborting Windows setup')
|
||||
|
||||
# Release currently used volume letters (ensures that the drives will get S, T, & W as needed below)
|
||||
remove_volume_letters(keep=windows_image['Source'])
|
||||
|
||||
# Format and partition drive
|
||||
print('\n Formatting Drive... \t\t', end='', flush=True)
|
||||
try:
|
||||
if (dest_disk['Use GPT']):
|
||||
format_gpt(dest_disk, windows_version['Family'])
|
||||
else:
|
||||
format_mbr(dest_disk, windows_version['Family'])
|
||||
print_success('Complete.')
|
||||
except:
|
||||
# We need to crash as the drive is in an unknown state
|
||||
print_error('Failed.')
|
||||
raise
|
||||
|
||||
# Apply Image
|
||||
print(' Applying Image... \t\t', end='', flush=True)
|
||||
try:
|
||||
setup_windows(bin, windows_image, windows_version)
|
||||
print_success('Complete.')
|
||||
except subprocess.CalledProcessError:
|
||||
print_error('Failed.')
|
||||
errors = True
|
||||
except:
|
||||
# We need to crash as the OS is in an unknown state
|
||||
print_error('Failed.')
|
||||
raise
|
||||
|
||||
# Create Boot files
|
||||
print(' Update Boot Partition...\t\t', end='', flush=True)
|
||||
try:
|
||||
update_boot_partition()
|
||||
print_success('Complete.')
|
||||
except subprocess.CalledProcessError:
|
||||
# Don't need to crash as this is (potentially) recoverable
|
||||
print_error('Failed.')
|
||||
errors = True
|
||||
except:
|
||||
print_error('Failed.')
|
||||
raise
|
||||
|
||||
# Setup WinRE
|
||||
print(' Update Recovery Tools...\t\t', end='', flush=True)
|
||||
try:
|
||||
setup_windows_re(windows_version)
|
||||
print_success('Complete.')
|
||||
except SetupError:
|
||||
print('Skipped.')
|
||||
except:
|
||||
# Don't need to crash as this is (potentially) recoverable
|
||||
print_error('Failed.')
|
||||
errors = True
|
||||
|
||||
# Print summary
|
||||
if errors:
|
||||
print_warning('\nErrors were encountered during setup.')
|
||||
time.sleep(30)
|
||||
else:
|
||||
print_success('\nNo errors were encountered during setup.')
|
||||
time.sleep(5)
|
||||
pause('\nPress Enter to return to main menu... ')
|
||||
|
||||
def menu_tools():
|
||||
tools = [
|
||||
{'Name': 'Blue Screen View', 'Folder': 'BlueScreenView', 'File': 'BlueScreenView.exe'},
|
||||
{'Name': 'Fast Copy', 'Folder': 'FastCopy', 'File': 'FastCopy.exe', 'Args': ['/log', '/logfile=X:\WK\Info\FastCopy.log', '/cmd=noexist_only', '/utf8', '/skip_empty_dir', '/linkdest', '/open_window', '/balloon=FALSE', r'/exclude=$RECYCLE.BIN;$Recycle.Bin;.AppleDB;.AppleDesktop;.AppleDouble;.com.apple.timemachine.supported;.dbfseventsd;.DocumentRevisions-V100*;.DS_Store;.fseventsd;.PKInstallSandboxManager;.Spotlight*;.SymAV*;.symSchedScanLockxz;.TemporaryItems;.Trash*;.vol;.VolumeIcon.icns;desktop.ini;Desktop?DB;Desktop?DF;hiberfil.sys;lost+found;Network?Trash?Folder;pagefile.sys;Recycled;RECYCLER;System?Volume?Information;Temporary?Items;Thumbs.db']},
|
||||
{'Name': 'HWiNFO', 'Folder': 'HWiNFO', 'File': 'HWiNFO.exe'},
|
||||
{'Name': 'NT Password Editor', 'Folder': 'NT Password Editor', 'File': 'ntpwedit.exe'},
|
||||
{'Name': 'Notepad++', 'Folder': 'NotepadPlusPlus', 'File': 'notepadplusplus.exe'},
|
||||
{'Name': 'PhotoRec', 'Folder': 'TestDisk', 'File': 'photorec_win.exe', 'Args': ['-new_console:n']},
|
||||
{'Name': 'Prime95', 'Folder': 'Prime95', 'File': 'prime95.exe'},
|
||||
{'Name': 'ProduKey', 'Folder': 'ProduKey', 'File': 'ProduKey.exe', 'Args': ['/external', '/ExtractEdition:1']},
|
||||
{'Name': 'Q-Dir', 'Folder': 'Q-Dir', 'File': 'Q-Dir.exe'},
|
||||
{'Name': 'TestDisk', 'Folder': 'TestDisk', 'File': 'testdisk_win.exe', 'Args': ['-new_console:n']},
|
||||
]
|
||||
actions = [
|
||||
{'Name': 'Main Menu', 'Letter': 'M'},
|
||||
]
|
||||
|
||||
# Menu loop
|
||||
while True:
|
||||
selection = menu_select('Tools Menu', tools, actions)
|
||||
|
||||
if (selection.isnumeric()):
|
||||
tool = tools[int(selection)-1]
|
||||
cmd = ['{bin}\\{folder}\\{file}'.format(bin=bin, folder=tool['Folder'], file=tool['File'])]
|
||||
if tool['Name'] == 'Blue Screen View':
|
||||
# Select path to scan
|
||||
minidump_path = select_minidump_path()
|
||||
if minidump_path is not None:
|
||||
tool['Args'] = ['/MiniDumpFolder', minidump_path]
|
||||
if 'Args' in tool:
|
||||
cmd += tool['Args']
|
||||
try:
|
||||
subprocess.Popen(cmd)
|
||||
except:
|
||||
print_error('Failed to run {prog}'.format(prog=tool['Name']))
|
||||
time.sleep(2)
|
||||
elif (selection == 'M'):
|
||||
break
|
||||
|
||||
def menu_main():
|
||||
menus = [
|
||||
{'Name': 'Create Backups', 'Menu': menu_backup_imaging},
|
||||
{'Name': 'Setup Windows', 'Menu': menu_windows_setup},
|
||||
{'Name': 'Misc Tools', 'Menu': menu_tools},
|
||||
]
|
||||
actions = [
|
||||
{'Name': 'Command Prompt', 'Letter': 'C'},
|
||||
{'Name': 'Reboot', 'Letter': 'R'},
|
||||
{'Name': 'Shutdown', 'Letter': 'S'},
|
||||
]
|
||||
|
||||
# Main loop
|
||||
while True:
|
||||
selection = menu_select('Main Menu', menus, actions, secret_exit=True)
|
||||
|
||||
if (selection.isnumeric()):
|
||||
try:
|
||||
menus[int(selection)-1]['Menu']()
|
||||
except AbortError:
|
||||
pass
|
||||
except:
|
||||
print_error('Major exception in: {menu}'.format(menu=menus[int(selection)-1]['Name']))
|
||||
print_warning(' Please let The Wizard know and he\'ll look into it (Please include the details below).')
|
||||
print(traceback.print_exc())
|
||||
time.sleep(5)
|
||||
print_info(' You can retry but if this crashes again an alternative approach may be required.')
|
||||
pause('\nPress enter to return to the main menu')
|
||||
elif (selection == 'C'):
|
||||
run_program(['cmd', '-new_console:n'], check=False)
|
||||
elif (selection == 'R'):
|
||||
run_program(['wpeutil', 'reboot'])
|
||||
elif (selection == 'S'):
|
||||
run_program(['wpeutil', 'shutdown'])
|
||||
else:
|
||||
sys.exit(0)
|
||||
|
||||
if __name__ == '__main__':
|
||||
menu_main()
|
||||
68
Scripts/settings/main.py
Normal file
68
Scripts/settings/main.py
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
# Wizard Kit PE: Settings - Main / Branding
|
||||
|
||||
# Features
|
||||
ENABLED_UPLOAD_DATA = False
|
||||
|
||||
# STATIC VARIABLES (also used by .cmd files)
|
||||
## Not using spaces aroung '=' for easier .cmd substrings
|
||||
ARCHIVE_PASSWORD='Abracadabra'
|
||||
KIT_NAME_FULL='Wizard Kit PE'
|
||||
KIT_NAME_SHORT='WKPE'
|
||||
OFFICE_SERVER_IP='10.0.0.10'
|
||||
QUICKBOOKS_SERVER_IP='10.0.0.10'
|
||||
SUPPORT_MESSAGE='Please let 2Shirt know by opening an issue on GitHub'
|
||||
TIME_ZONE='Pacific Standard Time' # Always use "Standard Time" (DST is applied correctly)
|
||||
|
||||
# SERVER VARIABLES
|
||||
## NOTE: Windows can only use one user per server. This means that if
|
||||
## one server serves multiple shares then you have to use the same
|
||||
## user/password for all of those shares.
|
||||
BACKUP_SERVERS = [
|
||||
{ 'IP': '10.0.0.10',
|
||||
'Name': 'ServerOne',
|
||||
'Mounted': False,
|
||||
'Share': 'Backups',
|
||||
'User': 'restore',
|
||||
'Pass': 'Abracadabra',
|
||||
},
|
||||
{ 'IP': '10.0.0.11',
|
||||
'Name': 'ServerTwo',
|
||||
'Mounted': False,
|
||||
'Share': 'Backups',
|
||||
'User': 'restore',
|
||||
'Pass': 'Abracadabra',
|
||||
},
|
||||
]
|
||||
CLIENT_INFO_SERVER = {
|
||||
'IP': '10.0.0.10',
|
||||
'RegEntry': r'0x10001,0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000',
|
||||
'Share': '/srv/ClientInfo',
|
||||
'User': 'upload',
|
||||
}
|
||||
OFFICE_SERVER = {
|
||||
'IP': OFFICE_SERVER_IP,
|
||||
'Name': 'ServerOne',
|
||||
'Mounted': False,
|
||||
'Share': 'Office',
|
||||
'User': 'restore',
|
||||
'Pass': 'Abracadabra',
|
||||
}
|
||||
QUICKBOOKS_SERVER = {
|
||||
'IP': QUICKBOOKS_SERVER_IP,
|
||||
'Name': 'ServerOne',
|
||||
'Mounted': False,
|
||||
'Share': 'QuickBooks',
|
||||
'User': 'restore',
|
||||
'Pass': 'Abracadabra',
|
||||
}
|
||||
WINDOWS_SERVER = {
|
||||
'IP': '10.0.0.10',
|
||||
'Name': 'ServerOne',
|
||||
'Mounted': False,
|
||||
'Share': 'Windows',
|
||||
'User': 'restore',
|
||||
'Pass': 'Abracadabra',
|
||||
}
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("This file is not meant to be called directly.")
|
||||
55
Scripts/settings/tools.py
Normal file
55
Scripts/settings/tools.py
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
# Wizard Kit PE: Settings - Tools
|
||||
|
||||
TOOLS = {
|
||||
# NOTE: BinDir will be prepended to these paths at runtime
|
||||
'AIDA64': {
|
||||
'32': r'AIDA64\aida64.exe'},
|
||||
'AutoRuns': {
|
||||
'32': r'Autoruns\autoruns.exe',
|
||||
'64': r'Autoruns\autoruns64.exe'},
|
||||
'BleachBit': {
|
||||
'32': r'BleachBit\bleachbit_console.exe'},
|
||||
'Caffeine': {
|
||||
'32': r'Caffeine\caffeine.exe'},
|
||||
'Du': {
|
||||
'32': r'Du\du.exe',
|
||||
'64': r'Du\du64.exe'},
|
||||
'ERUNT': {
|
||||
'32': r'ERUNT\ERUNT.EXE'},
|
||||
'Everything': {
|
||||
'32': r'Everything\Everything.exe',
|
||||
'64': r'Everything\Everything64.exe'},
|
||||
'FastCopy': {
|
||||
'32': r'FastCopy\FastCopy.exe',
|
||||
'64': r'FastCopy\FastCopy64.exe'},
|
||||
'HitmanPro': {
|
||||
'32': r'HitmanPro\HitmanPro.exe',
|
||||
'64': r'HitmanPro\HitmanPro64.exe'},
|
||||
'HWiNFO': {
|
||||
'32': r'HWiNFO\HWiNFO.exe',
|
||||
'64': r'HWiNFO\HWiNFO64.exe'},
|
||||
'KVRT': {
|
||||
'32': r'KVRT\KVRT.exe'},
|
||||
'NotepadPlusPlus': {
|
||||
'32': r'NotepadPlusPlus\notepadplusplus.exe'},
|
||||
'ProduKey': {
|
||||
'32': r'ProduKey\ProduKey.exe',
|
||||
'64': r'ProduKey\ProduKey64.exe'},
|
||||
'PuTTY-PSFTP': {
|
||||
'32': r'PuTTY\PSFTP.EXE'},
|
||||
'RKill': {
|
||||
'32': r'RKill\RKill.exe'},
|
||||
'SevenZip': {
|
||||
'32': r'7-Zip\7za.exe',
|
||||
'64': r'7-Zip\7za64.exe'},
|
||||
'TDSSKiller': {
|
||||
'32': r'TDSSKiller\TDSSKiller.exe'},
|
||||
'wimlib-imagex': {
|
||||
'32': r'wimlib\x32\wimlib-imagex.exe',
|
||||
'64': r'wimlib\x64\wimlib-imagex.exe'},
|
||||
'XMPlay': {
|
||||
'32': r'XMPlay\xmplay.exe'},
|
||||
}
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("This file is not meant to be called directly.")
|
||||
23
Scripts/winpe_root_menu.py
Normal file
23
Scripts/winpe_root_menu.py
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
# Wizard Kit PE: Root Menu
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Init
|
||||
os.chdir(os.path.dirname(os.path.realpath(__file__)))
|
||||
sys.path.append(os.getcwd())
|
||||
from functions.winpe_menus import *
|
||||
init_global_vars()
|
||||
set_title('{}: Root Menu'.format(KIT_NAME_FULL))
|
||||
global_vars['LogFile'] = r'{LogDir}\WinPE.log'.format(**global_vars)
|
||||
|
||||
if __name__ == '__main__':
|
||||
try:
|
||||
menu_root()
|
||||
except GenericAbort:
|
||||
# pause('Press Enter to return to main menu... ')
|
||||
pass
|
||||
except SystemExit:
|
||||
pass
|
||||
except:
|
||||
major_exception()
|
||||
|
|
@ -2,5 +2,5 @@
|
|||
[LaunchApps]
|
||||
wpeinit
|
||||
wpeutil updatebootinfo
|
||||
cd /d "%SystemDrive%\WK"
|
||||
"%SystemDrive%\WK\ConEmu\ConEmu.exe", /cmd cmd /k cd "%SystemDrive%\WK" & python "%SystemDrive%\WK\Scripts\menu.py"
|
||||
cd /d "%SystemDrive%\.bin"
|
||||
"%SystemDrive%\.bin\ConEmu\ConEmu.exe", /cmd cmd /k cd "%SystemDrive%\.bin" & python "%SystemDrive%\.bin\Scripts\menu.py"
|
||||
|
|
|
|||
|
|
@ -1,2 +1,2 @@
|
|||
@echo off
|
||||
python "%SystemDrive%\WK\Scripts\menu.py"
|
||||
python "%SystemDrive%\.bin\Scripts\menu.py"
|
||||
3
WK/_include/NotepadPlusPlus/npp.cmd
Normal file
3
WK/_include/NotepadPlusPlus/npp.cmd
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
@echo off
|
||||
|
||||
start "" %SystemDrive%\.bin\NotepadPlusPlus\notepadplusplus.exe %2 %3 %4 %5 %6 %7 %8 %9
|
||||
|
|
@ -1,33 +0,0 @@
|
|||
'// DISCLAIMER
|
||||
'// THIS COMES WITH NO WARRANTY, IMPLIED OR OTHERWISE. USE AT YOUR OWN RISK
|
||||
'// IF YOU ARE NOT COMFORTABLE EDITING THE REGISTRY THEN DO NOT USE THIS SCRIPT
|
||||
'//
|
||||
'// NOTES:
|
||||
'// This affects all users.
|
||||
'// This will prevent ANY executable named notepad.exe from running located anywhere on this computer!!
|
||||
'//
|
||||
'// Save this text to your notepad++ folder as a text file named npp.vbs (some AV don't like vbs, get a different AV :-P )
|
||||
'//
|
||||
'// USAGE
|
||||
'// 1)
|
||||
'// Navigate to registry key HKLM\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Image File Execution Options\
|
||||
'//
|
||||
' // 2)
|
||||
'// Add new subkey called notepad.exe
|
||||
'// This step is what tells windows to use the notepad++ exe, to undo simply delete this key
|
||||
'//
|
||||
'// 3)
|
||||
'// Create new Sting Value called Debugger
|
||||
'//
|
||||
'// 4)
|
||||
'// Modify value and enter wscript.exe "path to npp.vbs" e.g. wscript.exe "C:\Program Files\Notepad++\npp.vbs"
|
||||
|
||||
Option Explicit
|
||||
Dim sCmd, x
|
||||
sCmd = """" & LeftB(WScript.ScriptFullName, LenB(WScript.ScriptFullName) - LenB(WScript.ScriptName)) & "notepad++.exe" & """ """
|
||||
For x = 1 To WScript.Arguments.Count - 1
|
||||
sCmd = sCmd & WScript.Arguments(x) & " "
|
||||
Next
|
||||
sCmd = sCmd & """"
|
||||
CreateObject("WScript.Shell").Run sCmd, 1, True
|
||||
WScript.Quit
|
||||
Loading…
Reference in a new issue