Automating Windows & WSL Tasks with Python in VS Code

import subprocess
import os
from colorama import Fore
import platform
from util import display_text_color, display_alarm_color
import shutil
import sys
import ctypes
from pathlib import Path

def run_wsl_command(command, search_phrase=None):
    """
    This function runs a specified WSL command and prints the output.
    If a search phrase is provided, it filters the command output based on that phrase.
    Args:
        command (str): The WSL command to run
        search_phrase (str): Optional search phrase to filter command output. Defaults to None.
    Returns:
        None
    """
    if not command:
        # Display an error message if no command is provided
        display_alarm_color("No command provided to run in WSL.")
        return
    try:
        # Run the WSL command
        result = subprocess.run(
            # If a search phrase is provided, add it to the command
            ["wsl", command, search_phrase] if search_phrase else ["wsl", command],
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            check=True
        )
        # Display the command output
        display_text_color(result.stdout)
    except subprocess.CalledProcessError as e:
        # Display an error message if the command fails
        display_alarm_color(f"Command failed with error: {e}")
    except Exception as e:
        # Display an error message for any other exceptions
        display_alarm_color(f"An unexpected error occurred: {e}")

def run_help_command(search_phrase=None):
    """
    Execute the cmd.exe help command with an optional search phrase.
    Note: The Windows help command only lists available commands; it doesn't filter by phrase.

    Args:
        search_phrase (str, optional): A phrase to filter help output. Defaults to None.

    Returns:
        None
    """
    try:
        # Run the cmd.exe help command
        if search_phrase:
            # If a search phrase is provided, filter the help output
            cmd = ["cmd", "/c", f"help | findstr {search_phrase}"]
            # Run the filtered help command
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                encoding="utf-8",
                errors="replace",
                check=False # Do not raise exception on non-zero codes
            )
            if result.returncode == 0 and result.stdout.strip():
                # Display the filtered help output if matches are found
                display_text_color(f"Help for '{search_phrase}':", Fore.BLACK)
                display_text_color(result.stdout)
            else:
                # Display an informative message if no matches are found
                display_alarm_color(f'No help entry found matching "{search_phrase}".', Fore.YELLOW)
        else:
            # If no search phrase is provided, display the full help
            result = subprocess.run(
                ["cmd", "/c", "help"],
                capture_output=True,
                text=True,
                encoding="utf-8",
                errors="replace",
                check=True
            )
        display_text_color(result.stdout)
    except subprocess.CalledProcessError as e:
        # Display an error message if the command fails
        display_alarm_color(f"Command failed with error: {e}")
    except Exception as e:
        # Display an error message for any other exceptions
        display_alarm_color(f"An unexpected error occurred: {e}")

def is_admin():
    """
    Check if the script is running with administrative privileges.

    Returns:
        bool: True if running as an administrator, False otherwise.
    """
    try:
        # For non-Windows (Linux/macOS), check if UID is 0 (root)
        if platform.system() != "Windows":
            # Check if UID is 0
            return os.getuid() == 0
        else:
            # For Windows, use ctypes to check if user is admin and return True/False
            return ctypes.windll.shell32.IsUserAnAdmin() != 0
    except AttributeError:
        # This catch is typically for older Python versions on Windows that might
        # not have os.getuid(), or if there's an issue with ctypes.
        # It falls back to just the ctypes check for Windows.
        return ctypes.windll.shell32.IsUserAnAdmin() != 0

def elevate_privileges():
    """
    Attempt to re-run the current script with administrative privileges on Windows.

    This function will exit the current process and launch a new elevated one.

    Returns:
        bool: True if elevation was successful, False otherwise.
    """
    if platform.system() == "Windows":
        # Get the path to the current Python executable
        python_exe = sys.executable
        # Get the path to the current script
        script_path = os.path.abspath(sys.argv[0])

        # Prepare arguments to pass to the elevated process
        # The first argument is the script itself
        # sys.argv[1:] are any additional command-line arguments passed to the script
        arguments = [script_path] + sys.argv[1:]

        # Use ShellExecuteW with 'runas' verb to elevate privileges
        # The 'runas' verb triggers the UAC (User Account Control) prompt
        try:
            # shell32.ShellExecuteW(hwnd, lpOperation, lpFile, lpParameters, lpDirectory, nShowCmd)
            # hwnd: A handle to the parent window, 0 for no parent.
            # lpOperation: "runas" for elevation.
            # lpFile: The application to run (Python executable).
            # lpParameters: Arguments for the application (script path + its args).
            # lpDirectory: Working directory (current directory), None for current directory.
            # nShowCmd: How the window should be shown (SW_SHOWNORMAL = 1).
            ctypes.windll.shell32.ShellExecuteW(
                0, "runas", python_exe, " ".join(arguments), None, 1
            )
            # If ShellExecuteW succeeds, the current process will typically be exited
            # because the new elevated process takes over.
            sys.exit(0) # Exit the non-elevated process
        except Exception as e:
            # Display an error message if elevation fails
            display_alarm_color(f"Error attempting to elevate privileges: {e}")
            sys.exit(1) # Exit with an error code
            return False
    else:
        # Display an informative message for non-Windows (Linux/macOS)
        display_alarm_color("Privilege elevation is primarily a Windows-specific mechanism via UAC.")
        display_alarm_color("On Linux/macOS, you typically run with 'sudo' directly.")
        return False
    # If elevation is successful, this line won't be reached
    return True

def Windows_Maintenance():
    """
    Performs various Windows-specific maintenance tasks:
    - Empties the Recycle Bin.
    - Clears user and system temporary files.
    - Runs Disk Cleanup.
    - Initiates a Windows Defender full scan.

    Requires administrator privileges.

    Returns:
        None
    """
    if platform.system() != "Windows":
        # Display an error message if the OS is not Windows
        display_alarm_color("These tasks are Windows-specific and cannot be performed on this OS.")
        return

    if not is_admin():
        display_alarm_color("Administrator privileges required for Windows maintenance tasks. Attempting elevation...")
        # elevate_privileges() will exit the current script and relaunch an elevated one
        # The new elevated script instance will then re-enter this function and proceed.
        elevate_privileges()
        # This part of the code will only be reached if elevation fails,
        # or if the current process is not elevated.
        # In the context of `elevate_privileges` causing exit, this line won't be hit.
        return # Exit the non-elevated process if it somehow continues

    # --- Task 1: Empty the Windows Recycle Bin using cmd. ---
    display_text_color("\n--- Emptying Recycle Bin ---")
    try:
        # Run the cmd.exe command to empty the Recycle Bin
        # /c: Run the command
        # /rd: Delete the Recycle Bin
        # /s: Recursively delete subfolders
        # /q: Suppress error messages
        # C:\\$Recycle.Bin: Path to the Recycle Bin
        subprocess.run(["cmd.exe", "/c", "rd", "/s", "/q", "C:\\$Recycle.Bin"], check=True, capture_output=True, text=True)
        display_text_color("Recycle Bin emptied successfully. ✅")
    except subprocess.CalledProcessError as e:
        # Display an error message if the command fails
        display_alarm_color(f"Failed to empty Recycle Bin. Error: {e.stderr.strip()} ❌")

    # --- Task 2: Clear Temporary Files ---
    display_text_color("\n--- Clearing Temporary Files ---")
    temp_dirs = [
        Path(os.getenv("TEMP", "C:\\Windows\\Temp")), # User TEMP
        Path(os.getenv("WINDIR", "C:\\Windows")) / "Temp", # System Temp
        Path(os.getenv("LOCALAPPDATA", "")) / "Temp" # Another common user temp path
    ]

    for temp_dir in temp_dirs:
        # Loop through the list of temporary directories
        if not temp_dir.is_dir(): # Check if it's actually a directory
            display_alarm_color(f"Directory not found or not a directory: {temp_dir}. Skipping. ⚠️")
            continue

        display_text_color(f"Cleaning: {temp_dir}")
        for item in temp_dir.rglob("*"): # rglob is recursive, potentially slow for very large dirs
            try:
                # Try to delete the item
                if item.is_file():
                    # If it's a file, delete it
                    item.unlink()
                elif item.is_dir() and not list(item.iterdir()): # Only try to remove empty directories first
                    item.rmdir()
            except (PermissionError, OSError) as e:
                # Be more specific with error messages
                display_alarm_color(f"Could not delete {item}: {e.strerror or e}. Skipping. ⚠️")
            except Exception as e:
                # Display an error message if something unexpected happens
                display_alarm_color(f"An unexpected error occurred while deleting {item}: {e}. Skipping. ⚠️")

        # After individual item deletion attempts, try to remove remaining empty directories
        # This is a bit redundant if rmdir() worked above, but helps catch any missed empty ones.
        for item in temp_dir.iterdir():
            # Only try to remove empty directories
            if item.is_dir() and not list(item.iterdir()):
                try:
                    item.rmdir()
                except (PermissionError, OSError) as e:
                    display_alarm_color(f"Could not delete empty dir {item}: {e.strerror or e}. Skipping. ⚠️")

    display_text_color("Temporary files cleanup completed. ✅")

    # --- Task 3: Run Disk Cleanup ---
    display_text_color("\n--- Running Disk Cleanup ---")
    # cleanmgr /sagerun:1 executes a pre-configured cleanup (e.g., from cleanmgr /sageset:1). It refers to how Windows' Disk Cleanup tool (cleanmbgr) works.
    # The sageset and sagerun options are used to automate disk cleanup tasks by saving and executing a specific configuration profile.
    # Ensure a proper sageset configuration (e.g., 1) is saved before running this.
    # Press Win + S, type cmd, right-click “Command Prompt,” and select “Run as administrator.”
    # In the Command Prompt, type: cleanmgr /sageset:1.
    # Check the boxes for the file types you want cleanmgr /sagerun:1 to clean automatically.
    # Click “OK” to save the settings to profile ID 1 in the Windows registry
    try:
        subprocess.run(["cleanmgr", "/sagerun:1"], check=True, capture_output=True, text=True)
        display_text_color("Disk Cleanup initiated successfully. ✅ (Check system for full results)")
    except subprocess.CalledProcessError as e:
        display_alarm_color(f"Failed to initiate Disk Cleanup. Error: {e.stderr.strip()} ❌")

    # --- Task 4: Run Windows Defender Scan ---
    display_text_color("\n--- Running Windows Defender Scan ---")
    # Use environment variable to find Program Files for robustness
    program_files = os.getenv("ProgramFiles", r"C:\Program Files")
    # Construct the full path to the Windows Defender executable
    # C:\Program Files\Windows Defender\MpCmdRun.exe
    defender_path = Path(program_files) / "Windows Defender" / "MpCmdRun.exe"

    if not defender_path.exists():
        # Display an error message if the Windows Defender executable is not found
        display_alarm_color(f"Windows Defender executable not found at {defender_path}. Skipping scan. ⚠️")
        return

    # ScanType 2 = Full Scan
    # ScanType 1 = Quick Scan
    SCAN_TYPE_FULL = "2"
    SCAN_TYPE_QUICK = "1"

    # Pass command and arguments as a list for better security
    scan_cmd_list = [
        str(defender_path), # Convert Path object to string
        "-Scan",
        "-ScanType",
        SCAN_TYPE_FULL # Use named constant for clarity
    ]

    try:
        display_text_color("Starting Windows Defender full scan... This may take a while. ⏳")
        # Run the Windows Defender scan
        result = subprocess.run(scan_cmd_list, capture_output=True, text=True, check=True)
        display_text_color(f"Malware scan completed successfully. ✅\n{result.stdout}")
        if result.stderr:
            # Display any warnings or errors
            display_alarm_color(f"Malware scan warnings/errors: \n{result.stderr} ⚠️")
    except subprocess.CalledProcessError as e:
        # Display an error message if the scan fails
        display_alarm_color(f"Malware scan failed. Return code: {e.returncode}. Error:\n{e.stderr.strip()} ❌")
    except FileNotFoundError:
        # Display an error message if the Windows Defender executable is not found
        display_alarm_color(f"Windows Defender executable not found at '{defender_path}'. Please ensure Windows Defender is installed and accessible. ❌")
    except Exception as e:
        # Display an error message if something unexpected happens
        display_alarm_color(f"An unexpected error occurred during Windows Defender scan: {e} ❌")

    display_text_color("\n--- Windows Maintenance Complete ---")

# Example usage of the run_wsl_command function
if __name__ == "__main__":
    search_phrase = "mkdir"  # Optional search phrase for the command
    # You can call the function with or without a search phrase
    # run_wsl_command(search_phrase), e.g., run_wsl_command("ls")
    run_help_command("copy")
    # Call the Windows maintenance function to perform Windows-specific maintenance tasks
    Windows_Maintenance()