wizzardofOS_Windows.py

#Auth : ***
#Re-coded by Qyfashae 01/09/2023

import glob
import math
import os
import subprocess
import sys
import time
import traceback

# Define the sleep interval in seconds
SLEEP = 5

# Define the list of drive letters to choose from
LETTERS = 'BCDEFGHIJKLMNOPQRSTUVWXYZ'

# Define constants for menu options
REFRESH_LIST = 'REFRESH LIST'
REPLAY_DRIVE = 'D:\\'

# Define the main header text
MAIN_HEADER = """
ReplayWizard 01/09/2023

LEAVE THIS WINDOW OPEN WHILE YOU WORK!
When the op is complete, close all DSZ windows and choose QUIT.
"""

# Define the main function
def main():
    # Display the main header text
    print(MAIN_HEADER)

    # Initialize variables
    first_run = True
    sentinel_children = []

    # Start an infinite loop for menu selection
    while True:
        # Generate a list of available drives with DSZOpsDisk*
        drive_menu = [f'{x}:' for x in LETTERS if glob.glob(f'{x}:/DSZOpsDisk*')]

        # Check if there's only one drive and select it automatically
        if first_run and len(drive_menu) == 1:
            choice = drive_menu[0]
            print(f'Only one share found, selecting {choice} drive...')
            first_run = False
        else:
            # Display the drive menu and prompt for selection
            choice = menu(drive_menu, quitmsg='QUIT', text="Select the drive or share with the Windows OPS disk you'd like to use:")

        # Check if the user chose to quit
        if choice is None:
            # Terminate sentinel subprocesses and exit
            for child in sentinel_children:
                print(f'Stopping sentinel process: {child.pid}')
                child.terminate()
            break

        # Check if the user chose to refresh the list
        if choice == REFRESH_LIST:
            continue

        # Determine the path to the selected ops disk
        ops_disk_path = glob.glob(f'{choice}/DSZOpsDisk*')[0]
        replay_disk_path = os.path.join(REPLAY_DRIVE, 'ReplayDisk')

        # Check if the replay disk directory doesn't exist, and create it
        if not os.path.exists(replay_disk_path):
            create_replay_disk(ops_disk_path, replay_disk_path)

        # Determine the path to the logs directory
        logs_path = os.path.join(os.path.dirname(ops_disk_path), 'Logs')

        # Check if the logs directory exists
        if not os.path.exists(logs_path):
            print(f"\nCouldn't find the logs dir at {logs_path}")
            print('\nTry again when the WinOP DSZ GUI has launched.\n')
            continue

        # Determine the output directory for sentinel prompts
        output_dir = os.path.dirname(replay_disk_path)

        # Start the sentinel process and store it in the list
        sentinel_child = sentinel_prompts(ops_disk_path, output_dir)
        sentinel_children.append(sentinel_child)

        # Start the project menu for selected drive and replay disk
        project_menu(ops_disk_path, replay_disk_path)

# Define a function to create a replay disk from an ops disk
def create_replay_disk(ops_disk_path, replay_disk_path):
    # Print a message indicating the copying process
    print(f'\nCopying replay from {ops_disk_path} (may take a moment)...')

    # Change the current directory to the ops disk path
    os.chdir(ops_disk_path)

    # Create the path to the CreateReplay.py script
    script_path = os.path.join(ops_disk_path, 'CreateReplay.py')

    # Execute the CreateReplay.py script as a subprocess
    proc = subprocess.Popen([sys.executable, script_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
    
    # Communicate with the subprocess and provide the replay disk path
    proc.communicate(f'{replay_disk_path}\n')

# Define a function to handle sentinel prompts
def sentinel_prompts(ops_disk_path, output_dir):
    # Initialize flags for options
    get_files = False
    host_map = False

    # Check if the user wants to automatically rename GetFiles
    if yn_prompt('\nAutomatically rename GetFiles to closely match on-target names?'):
        print(f'\nFiles will save to {output_dir}GetFiles_Renamed')
        get_files = True

    # Check if the user wants to automatically aggregate network info
    if yn_prompt('\nAutomatically aggregate network info across targets?'):
        print(f'\nFiles will save to {output_dir}hostmap.txt')
        host_map = True

    # Check if neither option is selected, return if so
    if not (get_files or host_map):
        print()
        return

    # Define the path to the sentinel's fsmon.py script
    python_file = f'{ops_disk_path}/Resources/Ops/PyScripts/Windows/sentinel/fsmon.py'

    # Check if the sentinel script exists
    if not os.path.exists(python_file):
        print("\nCouldn't find sentinel! Skipping...")
        return

    # Construct the command line for running the sentinel
    command_line = [sys.executable, python_file, '--fresh', f'--output-dir={output_dir}']

    # Append additional options based on user choices
    if get_files:
        command_line.append('--get-files')
    if host_map:
        command_line.append('--host-table')

    # Define the path to the logs directory
    logs_path = os.path.join(os.path.dirname(ops_disk_path), 'Logs')

    # Append the logs path to the command line
    command_line.append(logs_path)

    # Print a message indicating the command that will be run
    print(f'\n\nRunning: {" ".join(command_line)}\n')

    # Start the sentinel as a subprocess and return the process object
    return subprocess.Popen(command_line)

# Define a function for the project menu
def project_menu(ops_disk_path, replay_disk_path):
    # Determine the path to the logs directory
    logs_path = os.path.join(os.path.dirname(ops_disk_path), 'Logs')

    # Initialize a flag for the first run
    first_run = True

    # Start an infinite loop for project selection
    while True:
        # List all directories in the logs directory as projects
        projects = [path for path in os.listdir(logs_path) if os.path.isdir(os.path.join(logs_path, path))]

        # Check if there's only one project and select it automatically
        if first_run and len(projects) == 1:
            choice = projects[0]
            print(f'Only one project found, auto-selecting {choice}...')
            first_run = False
        else:
            # Display the project menu and prompt for project selection
            choice = menu(projects, quitmsg='BACK TO SHARES LIST', text='Select a project to open the DSZ GUI:')

        # Check if the user chose to quit
        if choice is None:
            break

        # Check if the user chose to refresh the list
        if choice == REFRESH_LIST:
            continue

        # Determine the directory for the selected project's logs
        project_log_dir = os.path.join(logs_path, choice)

        # Write user defaults file for the selected project
        write_user_defaults_file(replay_disk_path, normpath(project_log_dir))

        # Print a message indicating the GUI launch
        print(f"\nLaunching GUI for '{choice}'... ", end='')

        # Construct the path to the script for launching the GUI
        script_path = os.path.join(replay_disk_path, 'start_lp.py')

        # Start the GUI as a subprocess
        subprocess.Popen([sys.executable, script_path])

        # Print a completion message
        print('done.\n')

        # Sleep for a specified duration before allowing another project selection
        print(f'Sleeping for {SLEEP} seconds, then you can select another project... ', end='')
        time.sleep(SLEEP)
        print('done.\n\n')

# Define a function to write user defaults file
def write_user_defaults_file(replay_disk_path, project_log_dir):
    # Determine the paths for configuration and resource directories
    config_dir = normpath(os.path.join(replay_disk_path, 'UserConfiguration'))
    resource_dir = normpath(os.path.join(replay_disk_path, 'Resources'))

    # Define the path for the user defaults file
    defaults_path = os.path.join(replay_disk_path, 'user.defaults')

    # Write user default settings to the file
    with open(defaults_path, 'w') as output:
        output.write('BuildType=true\n')
        output.write(f'LogDir={project_log_dir}\n')
        output.write('GuiType=true\n')
        output.write('wait.for.output=false\n')
        output.write('thread.dump=false\n')
        output.write(f'ConfigDir={config_dir}\n')
        output.write('OpMode=false\n ')
        output.write('LoadPrevious=true\n')
        output.write(f'ResourceDir={resource_dir}\n')
        output.write(f'OpsDisk={normpath(replay_disk_path)}\n')
        output.write('LocalAddress=00000001\n')
        output.write('LocalMode=false\n')
        output.write('replay.DSZ_KEYWORD=Extended')

# Define a function to normalize a file path for cross-platform compatibility
def normpath(path):
    # Use os.path.normpath to normalize the path
    norm_path = os.path.normpath(path)
    
    # Replace backslashes with double backslashes for proper escaping
    return norm_path.replace('\\', '\\\\')

# Define a menu function to display a menu and get user input
def menu(menu_list, text=None, quitmsg=None):
    # Add a "Refresh List" option to the menu
    menu_list = [REFRESH_LIST] + menu_list
    
    # Start an infinite loop for menu selection
    while True:
        # Calculate the number of digits needed for menu item alignment
        optspace = int(math.log(len(menu_list), 10))
        
        # Display an optional text message before the menu
        if text:
            print(f'\n{text}')
        
        # Display a quit message if provided
        if quitmsg:
            print(f'{optspace}d. {quitmsg}'.format(0))
        else:
            print('0. Quit')
        
        # Display the menu items with their corresponding numbers
        items = 0
        for i in menu_list:
            items += 1
            print(f'{optspace}d. {i}'.format(items))
        
        # Initialize the result to None
        result = None
        
        # Continue to prompt for user input until a valid selection is made
        while result is None or result > len(menu_list) or result < 0:
            result = prompt('Enter selection: ')
            try:
                result = int(result)
            except ValueError:
                result = None
            except TypeError:
                result = None
        
        # Check if the user chose to quit
        if result == 0:
            return None
        
        # Return the selected menu item from the list
        return menu_list[result - 1]

# Define a prompt function to display a message and get user input
def prompt(text, default=None):
    # Add the default value to the prompt if provided
    if default:
        text += f' [{default}] '
    
    # Get user input
    result = input(text)
    
    # Check if the user pressed Enter without input, return default if available
    if result == '':
        return None if default is None else str(default)
    else:
        return result

# Define a function to prompt for a yes/no response
def yn_prompt(text, default=True):
    # Create a prompt with a default value
    footer = ' [Y/n] > ' if default else ' [y/N] > '
    
    # Get user input, use 'y' as default if no input provided
    result = input(f'{text}{footer}').lower().strip()
    
    # If no input provided, return 'y' for default 'yes' response
    if not result:
        result = 'y' if default else 'n'
    
    # Return True if the response starts with 'n', indicating 'no'
    return result[0].find('n') < 0

# Check if the script is being run as the main program
if __name__ == '__main__':
    try:
        # Call the main function to start the script
        main()
    except:
        # Handle exceptions by printing an error message
        print()
        traceback.print_exc()
        print()
        print('An unrecoverable error has occurred. Cannot be wizardly for you.')
        print()
        
        # Wait for user input before exiting
        input('Press enter to exit.')