#To define a particular parameter, replace the 'parameterName' inside itsm.getParameter('parameterName') with that parameter's name

import os
import ctypes
from subprocess import PIPE, Popen


class disable_file_system_redirection:

    _disable = ctypes.windll.kernel32.Wow64DisableWow64FsRedirection
    _revert = ctypes.windll.kernel32.Wow64RevertWow64FsRedirection

    def __enter__(self):
        self.old_value = ctypes.c_long()
        self.success = self._disable(ctypes.byref(self.old_value))

    def __exit__(self, type, value, traceback):
        if self.success:
            self._revert(self.old_value)


# PERMANENT location - survives reboot and temp cleanup, readable at the lock screen.
Image_dir = r'C:\ProgramData\LockScreen'
Download_path = os.path.join(Image_dir, 'lockscreen.png')

url = itsm.getParameter('ImgURL')  # direct image URL, string datatype
# url = "https://www.comodo.com/new-assets/images/comodo-cybersecurity-managed-detection-and-response.png"

# --- Download script: force TLS 1.2, ensure folder exists, then download ---
Dwn_content = r'''
[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12
$dir = "''' + Image_dir + '''"
if(!(Test-Path $dir)) { New-Item -Path $dir -ItemType Directory -Force | Out-Null }
$url = "''' + url + '''"
$wc = New-Object System.Net.WebClient
$wc.DownloadFile($url, "''' + Download_path + '''")
'''

# --- Apply script: PersonalizationCSP (works on all editions) + policy enforcement ---
PsScript = r'''
$LockScreenImageValue = "''' + Download_path + r'''"

# 1) PersonalizationCSP - the mechanism the OS reads for the enforced lock screen
$CspPath = "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\PersonalizationCSP"
if(!(Test-Path $CspPath)) {
        New-Item -Path $CspPath -Force | Out-Null
}
New-ItemProperty -Path $CspPath -Name "LockScreenImageStatus" -Value 1 -PropertyType DWORD -Force | Out-Null
New-ItemProperty -Path $CspPath -Name "LockScreenImagePath"  -Value $LockScreenImageValue -PropertyType STRING -Force | Out-Null
New-ItemProperty -Path $CspPath -Name "LockScreenImageUrl"   -Value $LockScreenImageValue -PropertyType STRING -Force | Out-Null

# 2) Policy key - stops the user from changing it back (greys out the Settings option)
$PolPath = "HKLM:\SOFTWARE\Policies\Microsoft\Windows\Personalization"
if(!(Test-Path $PolPath)) {
        New-Item -Path $PolPath -Force | Out-Null
}
New-ItemProperty -Path $PolPath -Name "LockScreenImage"      -Value $LockScreenImageValue -PropertyType STRING -Force | Out-Null
New-ItemProperty -Path $PolPath -Name "NoChangingLockScreen" -Value 1 -PropertyType DWORD -Force | Out-Null

Stop-Process -processName Explorer -force -ErrorAction SilentlyContinue
'''


def ecmd(command):
    with disable_file_system_redirection():
        obj = Popen(command, shell=True, stdout=PIPE, stderr=PIPE)
    out, err = obj.communicate()
    return err


def CreateScriptFile(ps_content):
    try:
        file_name = 'ScriptFile.ps1'
        file_path = os.path.join(os.environ['TEMP'], file_name)
        with open(file_path, 'wb') as wr:
            wr.write(ps_content)
        return file_path
    except:
        return None


def RunPowerShellScript(ps_content):
    file_path = CreateScriptFile(ps_content)
    if file_path and os.path.exists(file_path):
        err = ecmd('powershell -NoProfile -ExecutionPolicy Bypass -File "%s"' % file_path)
        os.remove(file_path)
        if err == '':
            return True
        print "PowerShell stderr: " + err
        return False
    else:
        return False


# --- Run ---
RunPowerShellScript(Dwn_content)

if os.path.exists(Download_path):
    print "Downloaded successfully to " + Download_path
    if RunPowerShellScript(PsScript):
        print "Lock screen applied and persisted successfully"
    else:
        print "Script execution failed"
else:
    print "Download Failed"