using namespace Microsoft.Win32;

$null = New-Module {
    $wuPolicyPath = "HKLM:\SOFTWARE\Policies\Microsoft\Windows\WindowsUpdate";

    function Edit-DefaultUserKey {
        param(
            [scriptblock] $Action
        )

        $rootPath = "HKLM:\DefaultUser";
        $regRootPath = $rootPath -replace "^(\w+):([\\/]|$)","`$1`$2";
        $hivePath = "$env:SystemDrive\Users\Default\NTUSER.dat";
        $null = & reg load $regRootPath $hivePath;
        [RegistryKey] $root = Get-Item $rootPath;
        & $Action -Key $root;
        $root.Handle.Close();
        [System.GC]::Collect();
        & reg unload $regRootPath;
    }

    <#
        .SYNOPSIS
        Sets a message to show on the login screen.

        .PARAMETER Caption
        The title of the message.

        .PARAMETER Message
        The text of the message.
    #>
    function Set-BootMessage {
        param(
            [string] $Caption,
            [string] $Message
        )

        $options = @{
            legalnoticecaption = $Caption;
            legalnoticetext = $Message;
        };

        foreach ($key in $options.Keys) {
            Set-ItemProperty "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Policies\System" `
                -Name $key `
                -Type "String" `
                -Value ($options[$key]);
        }
    }

    <#
        .SYNOPSIS
        Disables the boot message.
    #>
    function Disable-BootMessage {
        Set-BootMessage;
    }

    <#
        .SYNOPSIS
        Disables automatic reboots by Windows Update.
    #>
    function Disable-WindowsUpdateAutoRestart {
        $path = "$wuPolicyPath\AU";
        $null = New-Item -Force $path -ErrorAction SilentlyContinue;

        Set-ItemProperty $path `
            -Name NoAutoRebootWithLoggedOnUsers `
            -Value 1;
    }

    <#
        .SYNOPSIS
        Resets the automatic reboot state.
    #>
    function Reset-WindowsUpdateAutoRestart {
        Remove-Item -Recurse "$wuPolicyPath";
    }
}