# Function to load a user's HKU registry hive
function UserReg {
    param (
        [string] $Username
    )
    # check if User folder exist
    if( -not (Test-Path -Path "C:\Users\$Username" -PathType Container) ) {
        Write-Host "User folder for $Username not found."
        Write-Host "You need to login to $Username in order for its files and registry to be setup."
        Write-Host "Press Enter to continue..."
        $null = Read-Host
        Write-Host "Continuing the script..."
    }
    Write-Host "Loading $Username's HKU registry hive..."
    $UserProfiles = Get-WmiObject Win32_UserProfile | Where-Object { $_.Special -eq $false }
    $UserProfile = $UserProfiles | Where-Object { $_.LocalPath.EndsWith("\$Username") }
    if ($null -ne $UserProfile) {
        $UserSID = $UserProfile.SID
        if( -not (Test-Path -Path "Registry::HKEY_USERS\$UserSID" -PathType Container) ) {
            Write-Host "REG LOAD HKEY_USERS\$UserSID C:\Users\$Username\NTUSER.DAT"
            $res = REG LOAD HKEY_USERS\$UserSID C:\Users\$Username\NTUSER.DAT
            if ($res -eq 1) {
                Write-Host "Failed to load $Username's HKU registry hive."
                return $null, $null
            }
        }
        Write-Host "Loaded $Username's HKU registry hive."
        return $UserSID, "Registry::HKEY_USERS\$UserSID"
    } else {
        Write-Host "User profile for $Username not found."
        return $null, $null
    }
}

function SetRegistry {
    param (
        [string] $regpath,
        [string] $regproperty
    )
    
    if( -not (Test-Path -Path $regpath -PathType Container) ) {
        New-Item -Path $regpath -Force -ItemType Registry
        New-ItemProperty -Path $regpath -Name $regproperty -Value 1 -PropertyType DWord
    }
    else {
        $RegistryItem = Get-ItemProperty -Path $regpath
        if ($RegistryItem.PSObject.Properties.Name -contains $regproperty) {
            Set-ItemProperty -Path $regpath -Name $regproperty -Value 1
        } else {
            New-ItemProperty -Path $regpath -Name $regproperty -Value 1 -PropertyType DWord
        }
    }
}

function SetupUserReg {
    param (
        [string] $username
    )
    $UserSID, $UserHKUPath = UserReg -Username $username

    if ($null -ne $UserSID -and $null -ne $UserHKUPath) {
        # Restrict access to Settings
        Write-Host "Restricting access to Settings..."
        $ControlPanelKeyPath = "$UserHKUPath\Software\Microsoft\Windows\CurrentVersion\Policies\Explorer"
        $ControlPanelValueName = "NoControlPanel"
        SetRegistry -regpath $ControlPanelKeyPath -regproperty $ControlPanelValueName
    
        # Disable access to regedit
        Write-Host "Disabling access to regedit..."
        $REGKeyPath = "$UserHKUPath\Software\Microsoft\Windows\CurrentVersion\Policies\System"
        $REGValueName = "DisableRegistryTools"
        SetRegistry -regpath $REGKeyPath -regproperty $REGValueName
    
        # Restrict access to Command Prompt
        Write-Host "Restricting access to Command Prompt..."
        $CMDKeyPath = "$UserHKUPath\Software\Policies\Microsoft\Windows\System"
        $CMDValueName = "DisableCMD"
        SetRegistry -regpath $CMDKeyPath -regproperty $CMDValueName
    
        # Add entries to DisallowRun for cmd.exe and powershell.exe
        Write-Host "Adding entries to DisallowRun for cmd.exe and powershell.exe..."
        $DisallowRunKeyPath = "$UserHKUPath\Software\Microsoft\Windows\CurrentVersion\Policies\Explorer"
        $DisallowRunValueName = "DisallowRun"
        $DisallowRunPath = "$UserHKUPath\Software\Microsoft\Windows\CurrentVersion\Policies\Explorer\DisallowRun"
        if (-not (Test-Path -Path "Registry::$DisallowRunPath" -PathType Container)) {
            New-Item -Path $DisallowRunPath -Force
        }
        Set-ItemProperty -Path $DisallowRunKeyPath -Name $DisallowRunValueName -Value 1
        $applications = @("cmd.exe", "powershell.exe", "powershell_ise.exe")
        Get-ItemProperty -Path $DisallowRunPath | ForEach-Object {
            Remove-ItemProperty -Path $DisallowRunPath -Name $_.PSObject.Properties.Name -ErrorAction SilentlyContinue
        }
        foreach ($valueName in $applications) {
            New-ItemProperty -Path $DisallowRunPath -Name $valueName -Value $valueName -PropertyType String
        }
        Write-Host "Added registry entries"
        Write-Host "Unloading registry HIVE"
        Write-Host "REG UNLOAD HKEY_USERS\$UserSID"
        $res = REG UNLOAD HKEY_USERS\$UserSID
        if ($res -eq 1) {
            Write-Host "Failed to unload $username's HKU registry hive."
        }
    } else {
        Write-Host "Unable to get the user's HKU registry."
    }
}