#Requires -RunAsAdministrator

$ErrorActionPreference = "Stop"

$env:SystemDirectory = [Environment]::SystemDirectory
$appCmd = $env:SystemDirectory + "\inetsrv\appcmd.exe"

function Test-IsAdmin {
    $currentUser = [Security.Principal.WindowsIdentity]::GetCurrent()
    $currentPrincipal = New-Object Security.Principal.WindowsPrincipal($currentUser)
    return $currentPrincipal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Prepare-Install-Directory([string]$InstallDir) {
	#If path is already exist returning
    if (Test-Path $InstallDir) {
		return
		#Remove-Item -LiteralPath $InstallDir -Force -Recurse
	}
    New-Item -ItemType Directory -Force -Path $InstallDir | Out-Null
    Write-Output "Install directory prepared successfully."
}

function Get-CLIInstallDir-From-InstallDir([string]$InstallDir) {
    $dir = "Appsentinels IISAgent"

    if ($InstallDir -eq "<auto>") {
        return (Join-Path $Env:ProgramFiles $dir)`
    }
    elseif (Test-Path $InstallDir -IsValid) {
        return $InstallDir
    }

    Write-Error "Invalid install directory provided '$InstallDir'"
}


function Install-Agent() {
    param(
        [string]$InstallDir = "<auto>",
        [Parameter(Mandatory = $false)]
        [switch]$Force
    )

    if (-not [Environment]::Is64BitProcess) {
        Write-Error "Install-Agent must be run from a 64-bit PowerShell session. Please relaunch using 64-bit PowerShell."
    }

    if (-not (Test-IsAdmin)) {
        Write-Error "This script must be run as an administrator."
    }

    if (-not $Force -and (Test-AgentInstallation)) {
        Write-Error "Run Uninstall-Agent first or use -Force to upgrade."
    }

    if ($Force -and (Test-AgentInstallation)) { Uninstall-Agent -Force }

    $installDir = Get-CLIInstallDir-From-InstallDir $InstallDir
    $dlPath = $PSScriptRoot
	Write-Output "dlPath: $dlPath"
    Prepare-Install-Directory -InstallDir $installDir

    try {
        # Install every bitness present in the package. A package may legitimately
        # ship only x64, in which case 32-bit pools stay uninstrumented and are
        # reported by Test-AgentBitnessCoverage during Register-Agent.
        $installed = @()

        foreach ($variant in Get-AgentModuleVariants) {
            $src = Join-Path $dlPath $variant.File
            if (-not (Test-Path $src)) {
                Write-Warning "No $($variant.File) in package. Application pools matching $($variant.PreCondition) will not be instrumented."
                continue
            }

            $machine = Get-DllMachineType $src
            if ($machine -ne $variant.Machine) {
                Write-Error "'$src' has PE machine type 0x$('{0:X}' -f $machine), which does not match $($variant.PreCondition). Rebuild it for the correct platform."
            }

            Copy-Item $src -Destination (Join-Path $installDir $variant.File) -Force
            $installed += $variant.File
        }

        if ($installed.Count -eq 0) {
            Write-Error "No module binaries found in '$dlPath'."
        }

        [System.Environment]::SetEnvironmentVariable('APPSENTINELS_AGENT_INSTALL_DIR', $installDir, [System.EnvironmentVariableTarget]::Machine)
        Write-Output "Installed Appsentinels IIS Agent Successfully ($($installed -join ', '))"
    }
    catch {
        $message = $_
        Write-Error "Could not install Appsentinels IIS Agent! $message"
    }
}

function Test-AgentInstallation {
	[System.Environment]::GetEnvironmentVariable('APPSENTINELS_AGENT_INSTALL_DIR', [System.EnvironmentVariableTarget]::Machine) -ne $null
}

function Test-AgentRegistration {
	foreach ($variant in Get-AgentModuleVariants) {
		if (Module-Exists $variant.Name) {
			return $true
		}
	}
	return $false
}

function Get-Temp-Directory() {
    $temp = $env:TEMP
    if (-not (Test-Path $temp)) {
        New-Item -ItemType Directory -Force -Path $temp | Out-Null
    }
    return $temp
}

function Get-Current-InstallDir() {
    $installDir = [System.Environment]::GetEnvironmentVariable("APPSENTINELS_AGENT_INSTALL_DIR", [System.EnvironmentVariableTarget]::Machine)

    if (-not $installDir -or -not (Test-Path $installDir)) {
        Write-Warning "APPSENTINELS_AGENT_INSTALL_DIR is not set or directory does not exist."
        return $null
    }
    return $installDir
}

function Module-Exists($moduleName) {
    $configPath = "system.webServer/globalModules"
    $module = Get-WebConfiguration -Filter "/$configPath/add[@name='$moduleName']"
    return $module -ne $null
}

function Get-ApplicationHostConfigPath {
    $path = "$($env:SystemDirectory)\inetsrv\config\applicationhost.config"

    if (-not (Test-Path -Path $path)) {
        Write-Error "The applicationhost.config file does not exist at the expected location: $path"
    }

    return $path
}

# The module is registered once per bitness. An application pool loads exactly one
# of these variants - IIS picks it by evaluating the preCondition against the worker
# process bitness - so a 32-bit (WOW64) pool is instrumented by the x86 build and a
# 64-bit pool by the x64 build. Names must be unique within globalModules.
function Get-AgentModuleVariants {
    return @(
        [PSCustomObject]@{ Name = 'AppsentinelsIISAgent';   PreCondition = 'bitness64'; File = 'AppsentinelsIISModule.dll';   Machine = 0x8664 }
        [PSCustomObject]@{ Name = 'AppsentinelsIISAgent32'; PreCondition = 'bitness32'; File = 'AppsentinelsIISModule32.dll'; Machine = 0x014C }
    )
}

# Reads the Machine field from the PE header so a mislabelled build is caught at
# install time rather than as a silent no-op in the worker process.
function Get-DllMachineType([string]$Path) {
    $fs = [IO.File]::OpenRead($Path)
    try {
        $br = New-Object IO.BinaryReader($fs)
        $fs.Position = 0x3C
        $fs.Position = $br.ReadInt32() + 4
        return $br.ReadUInt16()
    } finally {
        $fs.Close()
    }
}

function Register-Appsentinels-IISAgent-IIS([string]$InstallDir) {
    $installDir = Get-Current-InstallDir
    $configPath = "system.webServer/globalModules"
    $registered = @()

    foreach ($variant in Get-AgentModuleVariants) {
        $dll = Join-Path $installDir $variant.File

        if (-not (Test-Path $dll)) {
            Write-Warning "No module at '$dll'. Skipping $($variant.PreCondition) registration."
            continue
        }

        if (Module-Exists $variant.Name) {
            Write-Warning "$($variant.Name) module already exists. Overwriting"
            $selector = "/[name='$($variant.Name)',preCondition='$($variant.PreCondition)']"
            & $appCmd clear config -section:system.webServer/globalModules $selector /commit:apphost
        }

        Add-WebConfiguration -Filter "/$configPath" -Value @{name=$variant.Name; image=$dll; preCondition=$variant.PreCondition}
        $registered += "$($variant.Name) ($($variant.PreCondition))"
    }

    if ($registered.Count -eq 0) {
        Write-Error "No module binaries found under '$installDir'. Run Install-Agent first."
    }

    Write-Output "Appsentinels IIS Agent added as IIS native module: $($registered -join ', ')"
}

# Warns when 32-bit application pools exist that no installed variant can instrument.
# Without this the install reports success and those pools silently capture nothing.
function Test-AgentBitnessCoverage {
    Import-Module WebAdministration

    $installDir = Get-Current-InstallDir
    if (-not $installDir) {
        return
    }

    $pools32 = @(Get-ChildItem IIS:\AppPools | Where-Object { $_.enable32BitAppOnWin64 -eq $true })
    if ($pools32.Count -eq 0) {
        return
    }

    $x86Variant = Get-AgentModuleVariants | Where-Object { $_.PreCondition -eq 'bitness32' }
    $x86Dll = Join-Path $installDir $x86Variant.File

    if (Test-Path $x86Dll) {
        return
    }

    $names = ($pools32 | ForEach-Object { $_.Name }) -join ', '
    Write-Warning "These application pools run 32-bit (enable32BitAppOnWin64=True) but no x86 module is installed, so they will NOT be instrumented: $names"
}


function Register-Agent() {
     param(
        [string]$ControllerIP,
		[Parameter(Mandatory = $false)]
		[string]$ControllerPort = "9004",
		[Parameter(Mandatory = $false)]
		[string]$ControllerURI = "/mergedlog?sensor=IIS&cap=nohb"
    )
    if (-not (Test-IsAdmin)) {
        Write-Error "This script must be run as an administrator."
    }

    Import-Module WebAdministration

    if (-not (Test-AgentInstallation)) {
        Write-Error "Agent is not installed. Please run Install-Agent first."
    }

    $installDir = Get-Current-InstallDir
    if (-not $installDir) {
        Write-Error "Failed to determine the install directory."
    }

    Write-Output "Install-Directory:" $installDir

    $applicationHostConfigPath = Get-ApplicationHostConfigPath

    $tempDir = Get-Temp-Directory
    if (-not $tempDir) {
        Write-Error "Failed to get a temporary directory."
    } else {
        Copy-Item $applicationHostConfigPath -Destination $tempDir
        Write-Output "IIS config has been backed up to temp dir:" + $tempDir
    }

    if (-not $ControllerIP) {
        $ControllerIP = Read-Host -Prompt 'Enter Controller IP'
    }

    while (-not $ControllerIP) {
        Write-Error "Controller IP  cannot be empty."
        $ControllerIP = Read-Host -Prompt 'Enter Controller IP'
    }

    # Validate ControllerIP — must be a valid IPv4 address (e.g. 192.168.1.1) or FQDN (e.g. controller.example.com)
    $ipPattern   = '^(\d{1,3}\.){3}\d{1,3}$'
    $fqdnPattern = '^([a-zA-Z0-9]([a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}$'
    $isIPv4 = $ControllerIP -match $ipPattern
    $isFQDN = $ControllerIP -match $fqdnPattern

    if (-not $isIPv4 -and -not $isFQDN) {
        Write-Error "ControllerIP '$ControllerIP' is not a valid IPv4 address or FQDN."
        exit 1
    }

    if ($isIPv4) {
        foreach ($octet in $ControllerIP.Split('.')) {
            if ([int]$octet -gt 255) {
                Write-Error "ControllerIP '$ControllerIP' is not a valid IPv4 address. Octet '$octet' is out of range."
                exit 1
            }
        }
    }

    # Validate ControllerPort — must be numeric and in range 1-65535
    if ($ControllerPort -notmatch '^\d+$' -or [int]$ControllerPort -lt 1 -or [int]$ControllerPort -gt 65535) {
        Write-Error "ControllerPort '$ControllerPort' is not valid. Must be a number between 1 and 65535."
    }

    # Validate ControllerURI — must start with '/' and not exceed 512 chars
    if (-not $ControllerURI.StartsWith('/')) {
        Write-Error "ControllerURI '$ControllerURI' is not valid. It must start with '/'."
    }
    if ($ControllerURI.Length -gt 512) {
        Write-Error "ControllerURI is too long (max 512 characters). Got length: $($ControllerURI.Length)."
    }

	[System.Environment]::SetEnvironmentVariable('APPSENTINELS_CONTROLLER_IP', $ControllerIP, [System.EnvironmentVariableTarget]::Machine)
	[System.Environment]::SetEnvironmentVariable('APPSENTINELS_CONTROLLER_PORT', $ControllerPort, [System.EnvironmentVariableTarget]::Machine)
	[System.Environment]::SetEnvironmentVariable('APPSENTINELS_CONTROLLER_URI', $ControllerURI, [System.EnvironmentVariableTarget]::Machine)
	[System.Environment]::SetEnvironmentVariable('APPSENTINELS_LOGLEVEL', 'ERROR', [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_CAPTURE_UNKNOWN_SIZE_PAYLOAD', 'false', [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_ENABLE_CHUNKED_BODY_READ', 'false', [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_USE_SSL', '0', [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_SELF_SIGNED_CERT', '0', [System.EnvironmentVariableTarget]::Machine)


    Register-Appsentinels-IISAgent-IIS -InstallDir $installDir

    Test-AgentBitnessCoverage

    Reset-IIS
}


# Detects website- and application-scoped instrumentation. appcmd shows the
# EFFECTIVE config (inherited entries included), so scope is attributed by ruling
# out the broader level first: the scan is skipped entirely while global
# instrumentation exists, and a site hit skips the app scan beneath it (those apps
# merely inherit the site entry). Progress messages go to the host stream - anything
# sent to the pipeline here would corrupt the returned object.
function Check-If-Website-Instrumentation-Exists() {
    if (-Not (Test-Path $appCmd)) {
        Write-Error "appcmd.exe not found. Ensure that IIS is installed and appcmd.exe is in the expected location."
    }

    $instrumentedSites = @()
    $instrumentedApps = @()

    if (-not (Test-GlobalInstrumentation)) {
        $websites = & $appCmd list site /text:name
        foreach ($site in $websites) {
            $modules = & $appCmd list config "$site" /section:system.webServer/modules
            if ($modules -match 'AppsentinelsIISAgent') {
                Write-Host "Website '$site' is already instrumented with AppsentinelsIISAgent."
                $instrumentedSites += $site
                continue
            }

            foreach ($appPath in @(& $appCmd list app /site.name:"$site" /text:path)) {
                $rel = $appPath.Trim('/')
                if (-not $rel) { continue }
                $modules = & $appCmd list config "$site$appPath" /section:system.webServer/modules
                if ($modules -match 'AppsentinelsIISAgent') {
                    $rel = $rel -replace '/', '\'
                    Write-Host "Application '$site/$rel' is already instrumented with AppsentinelsIISAgent."
                    $instrumentedApps += [pscustomobject]@{ Website = $site; Application = $rel }
                }
            }
        }
    }

    return [pscustomobject]@{
        IsInstrumented = ($instrumentedSites.Count -gt 0 -or $instrumentedApps.Count -gt 0)
        Sites = $instrumentedSites
        Applications = $instrumentedApps
    }
}




function Add-GlobalInstrumentation() {
    param(
        [Parameter(Mandatory = $false)]
        [switch]$Force
    )

    if (-not (Test-IsAdmin)) {
        Write-Error "This script must be run as an administrator."
    }

    if (-not (Test-AgentInstallation)) {
        Write-Error "Agent is not installed. Please run Install-Agent first."
    }

    if (-not (Test-AgentRegistration)) {
        Write-Error "Agent is not registered. Please run Register-Agent first."
    }

    if (Test-GlobalInstrumentation) {
        Write-Output "Global Instrumentation already exists."
        return
    }

    # A site- or app-scoped <add> plus a global one duplicates the entry in that
    # subtree's effective config and 500.19s it, so narrower scopes must go first.
    $result = Check-If-Website-Instrumentation-Exists
    if ($result.IsInstrumented) {
        if (-not $Force) {
            Write-Error "Website/application instrumentation already exists. Remove it first or use -Force"
        }
        foreach ($site in $result.Sites) {
            Remove-AppInstrumentation -website $site
        }
        foreach ($app in $result.Applications) {
            Remove-AppInstrumentation -website $app.Website -application $app.Application
        }
    }

    $lockInfo = Get-WebConfigurationLock -Filter "//modules"
    if ($lockInfo) {
        if ($Force) {
            Remove-WebConfigurationLock -Filter "//modules" -Force
            Write-Output "Unlocked the modules section at the server level."
        } else {
            Write-Error "Modules section is locked for websites. Unlock it before instrumenting."
        }
    }
    # Only enable variants that are actually registered in globalModules - a modules
    # entry pointing at a missing global module makes IIS fail the config with 500.19.
    foreach ($variant in Get-AgentModuleVariants) {
        if (Module-Exists $variant.Name) {
            Update-WebModuleConfiguration -ModuleName $variant.Name -PreCondition $variant.PreCondition
        }
    }

    Write-Output "IIS is instrumented globally"
}

function Update-WebModuleConfiguration {
    param(
        [string]$ModuleName,
        [string]$PreCondition,
        [string]$Website,
        [string]$Application
    )

    if (-not $Website) {
        $psPath = 'MACHINE/WEBROOT/APPHOST'
    } elseif (-not $Application) {
        $psPath = "IIS:\Sites\$Website"
    } else {
        $psPath = "IIS:\Sites\$Website\$Application"
    }

    $filter = "//modules/add[@name='$ModuleName']"

    if ((Get-WebConfiguration -PSPath $psPath -Filter $filter -ErrorAction SilentlyContinue) -ne $null) {
        # Module configuration exists, update it
        Set-WebConfigurationProperty -PSPath $psPath -Filter $filter -Name "preCondition" -Value $PreCondition
        Write-Output "Configuration updated for $ModuleName with precondition $PreCondition"
    } else {
        # Module configuration does not exist, add it
        Add-WebConfiguration -PSPath $psPath -Filter "//modules" -Value @{name=$ModuleName; preCondition=$PreCondition}
        Write-Output "Configuration added for $ModuleName with precondition $PreCondition"
    }
}


function Reset-IIS() {
    Start-Process "iisreset.exe" -NoNewWindow -Wait
}


function Uninstall-Agent() {
    param(
        [Parameter(Mandatory = $false)]
        [switch]$Force
    )

    if (-not (Test-IsAdmin)) {
        Write-Error "This script must be run as an administrator."
    }

    if (-not (Test-AgentInstallation)) {
        Write-Output "Agent is not installed."
        return
    }

    if ($Force) {
        Unregister-Agent -Force
    }

    if (Test-AgentRegistration) {
        Write-Error "Agent is registered as native module. Run Unregister-Agent First or use -Force."
    }

    $installDir = Get-Current-InstallDir
    #Remove-Item -LiteralPath $installDir -Force -Recurse
    foreach ($variant in Get-AgentModuleVariants) {
        $dllPath = Join-Path $installDir $variant.File
        if (Test-Path $dllPath) {
            Remove-Item -LiteralPath $dllPath -Force
        }
    }

    if ((Test-Path $installDir) -and ((Get-ChildItem -LiteralPath $installDir -Force | Measure-Object).Count -eq 0)) {
         Remove-Item -LiteralPath $installDir -Force
    }

    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_AGENT_INSTALL_DIR', $null, [System.EnvironmentVariableTarget]::Machine)

    Write-Output "Appsentinels IIS Agent uninstall successful"
}

function Unregister-Agent() {
    param(
        [Parameter(Mandatory = $false)]
        [switch]$Force
    )

    if (-not (Test-IsAdmin)) {
        Write-Error "This script must be run as an administrator."
    }

    if (-not (Test-AgentRegistration)) {
        Write-Warning "Agent is not registered as native module."
        return
    }

    if ($Force) {
        Uninstrument-IIS
    } else {

        $result = Check-If-Website-Instrumentation-Exists
        if ($result.IsInstrumented) {
            Write-Error "Website/application instrumentation exists. Run Remove-AppInstrumentation first."
        }

        if (Test-GlobalInstrumentation) {
            Write-Error "Global Instrumentation Exists. Run Remove-GlobalInstrumentation First."
        }
    }

    if (-Not (Test-Path $appCmd)) {
        Write-Error "appcmd.exe not found. Ensure that IIS is installed and appcmd.exe is in the expected location."
    }

    # Every registered variant must go. A leftover entry pointing at a deleted DLL
    # stops worker processes of that bitness from starting at all.
    foreach ($variant in Get-AgentModuleVariants) {
        if (Module-Exists $variant.Name) {
            $selector = "/[name='$($variant.Name)',preCondition='$($variant.PreCondition)']"
            & $appCmd clear config -section:system.webServer/globalModules $selector /commit:apphost
        }
    }

    $env:dll = $null

    # Clean all the environment variables that we might set
    [System.Environment]::SetEnvironmentVariable("APPSENTINELS_CONTROLLER_IP", $null , [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable("APPSENTINELS_CONTROLLER_PORT", $null , [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable("APPSENTINELS_CONTROLLER_URI", $null , [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable("APPSENTINELS_TO_CONTROLLER_PROTOCOL_SCHEME", $null , [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_MAX_PAYLOAD_SIZE', $null  , [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_SUPPORTED_CONTENT_TYPE', $null  , [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_BYPASS_METHOD', $null  , [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_CAPTURE_UNKNOWN_SIZE_PAYLOAD', $null , [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_LOGLEVEL', $null  , [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_ENABLE_CHUNKED_BODY_READ', $null, [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_USE_SSL', $null, [System.EnvironmentVariableTarget]::Machine)
    [System.Environment]::SetEnvironmentVariable('APPSENTINELS_SELF_SIGNED_CERT', $null, [System.EnvironmentVariableTarget]::Machine)

    Write-Output "Removed IISAgent as native module from iis"
    Reset-IIS
}

function Uninstrument-IIS() {
    # Global first: with it gone, the scan below attributes the remaining entries to
    # their true site/app scope instead of seeing inherited copies everywhere.
    if(Test-GlobalInstrumentation) {
        Remove-GlobalInstrumentation
    }

    $result = Check-If-Website-Instrumentation-Exists
    foreach ($site in $result.Sites) {
        Remove-AppInstrumentation -website $site
    }
    foreach ($app in $result.Applications) {
        Remove-AppInstrumentation -website $app.Website -application $app.Application
    }
}

function Remove-GlobalInstrumentation() {
    if (-not (Test-IsAdmin)) {
        Write-Error "This script must be run as an administrator."
    }

    if (-not (Test-GlobalInstrumentation)) {
        Write-Output "Global instrumentation doesn't exist."
        return
    }

    foreach ($variant in Get-AgentModuleVariants) {
        Clear-WebConfiguration -PSPath 'MACHINE/WEBROOT/APPHOST' -Filter "//modules/add[@name='$($variant.Name)']" -ErrorAction SilentlyContinue
    }

    Write-Output "Global instrumentation is removed now."
}

function Test-GlobalInstrumentation {
    $configPath = "system.webServer/modules"

    foreach ($variant in Get-AgentModuleVariants) {
        $module = Get-WebConfiguration -PSPath "MACHINE/WEBROOT/APPHOST" -Filter "/$configPath/add[@name='$($variant.Name)']"
        if ($module -ne $null) {
            return $true
        }
    }

    return $false
}

# Returns $true when any agent module variant is enabled in the effective <modules>
# config at the given PSPath. Inherited entries count as hits, so callers must rule
# out broader scopes first (see the guard ordering in Add-AppInstrumentation).
function Test-InstrumentationAtPath([string]$PSPath) {
    foreach ($variant in Get-AgentModuleVariants) {
        if ((Get-WebConfiguration -PSPath $PSPath -Filter "//modules/add[@name='$($variant.Name)']" -ErrorAction SilentlyContinue) -ne $null) {
            return $true
        }
    }
    return $false
}

function Add-AppInstrumentation() {
    param(
        [Parameter(Mandatory = $true)]
        [string]$website,
        [Parameter(Mandatory = $false)]
        [string]$application,
        [Parameter(Mandatory = $false)]
        [switch]$Force
    )

    if (-not (Test-IsAdmin)) {
        Write-Error "This script must be run as an administrator."
    }

    if (-not (Test-AgentInstallation)) {
        Write-Error "Agent is not installed. Please run Install-Agent first."
    }

    if (-not (Test-AgentRegistration)) {
        Write-Error "Agent is not registered. Please run Register-Agent first."
    }

    if ((Test-GlobalInstrumentation) -and (-not $Force)) {
        Write-Error "Global instrumentation already exists. Remove it before instrumenting website or use -Force"
    } elseif ((Test-GlobalInstrumentation) -and ($Force)) {
        Remove-GlobalInstrumentation
    }

    if (-not (Test-Path "IIS:\Sites\$website")) {
        Write-Error "Website '$website' does not exist in IIS."
    }

    if ($application -and -not (Test-Path "IIS:\Sites\$website\$application")) {
        Write-Error "Application '$application' does not exist under website '$website'."
    }

    # Duplicate <add> entries for the same module at two levels of one inheritance
    # chain break the whole subtree with 500.19, so overlapping scopes are rejected.
    # Global instrumentation was ruled out (or removed via -Force) above, therefore a
    # hit at the site here can only be site-scoped; likewise a hit at an application
    # while the site is clean can only be app-scoped.
    if ($application) {
        if (Test-InstrumentationAtPath "IIS:\Sites\$website") {
            if ($Force) {
                Remove-AppInstrumentation -website $website
            } else {
                Write-Error "Website '$website' is already instrumented and application '$application' inherits it. Remove the website instrumentation first or use -Force."
            }
        }
    } elseif (-not (Test-InstrumentationAtPath "IIS:\Sites\$website")) {
        foreach ($app in @(Get-WebApplication -Site $website)) {
            $appPath = $app.Path.Trim('/') -replace '/', '\'
            if (-not $appPath) { continue }
            if (Test-InstrumentationAtPath "IIS:\Sites\$website\$appPath") {
                if ($Force) {
                    Remove-AppInstrumentation -website $website -application $appPath
                } else {
                    Write-Error "Application '$website/$appPath' is already instrumented and would conflict with website-level instrumentation. Remove it first or use -Force."
                }
            }
        }
    }

    $lockInfo = Get-WebConfigurationLock -Filter "//modules"
    if ($lockInfo) {
        if ($Force) {
            Remove-WebConfigurationLock -Filter "//modules" -Force
            Write-Output "Unlocked the modules section at the server level."
        } else {
            Write-Error "Modules section is locked for websites. Unlock it before instrumenting."
        }
    }

    $websiteLockInfo = Get-WebConfigurationLock -PSPath "IIS:\Sites\$website" -Filter "//modules"
    if ($websiteLockInfo) {
        if ($Force) {
            Remove-WebConfigurationLock -PSPath "IIS:\Sites\$website" -Filter "//modules" -Force
            Write-Output "Removed the lock from the modules section for website $website."
        } else {
            Write-Error "Modules section is locked for website $website. Unlock it before instrumenting."
        }
    }

    if ($application) {
        $appLockInfo = Get-WebConfigurationLock -PSPath "IIS:\Sites\$website\$application" -Filter "//modules"
        if ($appLockInfo) {
            if ($Force) {
                Remove-WebConfigurationLock -PSPath "IIS:\Sites\$website\$application" -Filter "//modules" -Force
                Write-Output "Removed the lock from the modules section for application $website/$application."
            } else {
                Write-Error "Modules section is locked for application $website/$application. Unlock it before instrumenting."
            }
        }
    }

    foreach ($variant in Get-AgentModuleVariants) {
        if (Module-Exists $variant.Name) {
            Update-WebModuleConfiguration -ModuleName $variant.Name -PreCondition $variant.PreCondition -Website $website -Application $application
        }
    }

    if ($application) {
        Write-Output "IIS is instrumented for application $website/$application"
    } else {
        Write-Output "IIS is instrumented for website $website"
    }
}

function Remove-AppInstrumentation() {
    param(
        [Parameter(Mandatory = $true)]
        [string]$website,
        [Parameter(Mandatory = $false)]
        [string]$application
    )

    if (-not (Test-IsAdmin)) {
        Write-Error "This script must be run as an administrator."
    }

    if ($application) {
        $psPath = "IIS:\Sites\$website\$application"
    } else {
        $psPath = "IIS:\Sites\$website"
    }

    foreach ($variant in Get-AgentModuleVariants) {
        Clear-WebConfiguration -PSPath $psPath -Filter "//modules/add[@name='$($variant.Name)']" -ErrorAction SilentlyContinue
    }

    if ($application) {
        Write-Output "IIS instrumentation is removed for application $website/$application"
    } else {
        Write-Output "IIS instrumentation is removed for website $website"
    }
}





<#
# Reference: default values for optional environment variables.
# Use Update-EnvironmentalVariable to set any new envs or update any existing envs variable.
# Update-EnvironmentalVariable -Name 'APPSENTINELS_TO_CONTROLLER_PROTOCOL_SCHEME' -Value 'http'
# Update-EnvironmentalVariable -Name 'APPSENTINELS_MAX_PAYLOAD_SIZE'               -Value '131072'
# Update-EnvironmentalVariable -Name 'APPSENTINELS_SUPPORTED_CONTENT_TYPE'         -Value 'json,xml,graphql,form'
# Update-EnvironmentalVariable -Name 'APPSENTINELS_BYPASS_METHOD'                  -Value 'OPTIONS,HEAD'
# Update-EnvironmentalVariable -Name 'APPSENTINELS_CAPTURE_UNKNOWN_SIZE_PAYLOAD'   -Value 'false'
# Update-EnvironmentalVariable -Name 'APPSENTINELS_ENABLE_CHUNKED_BODY_READ'       -Value 'false'
#>
function Update-EnvironmentalVariable() {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Name,
        [Parameter(Mandatory = $true)]
        [string]$Value
    )

    if (-not (Test-IsAdmin)) {
        Write-Error "This script must be run as an administrator."
    }

    [System.Environment]::SetEnvironmentVariable($Name, $Value, [System.EnvironmentVariableTarget]::Machine)
    Write-Output "Environment variable '$Name' set successfully."
    Reset-IIS
}


Export-ModuleMember -Function Install-Agent
Export-ModuleMember -Function Register-Agent
Export-ModuleMember -Function Update-EnvironmentalVariable
Export-ModuleMember -Function Add-GlobalInstrumentation
Export-ModuleMember -Function Remove-GlobalInstrumentation
Export-ModuleMember -Function Unregister-Agent
Export-ModuleMember -Function Uninstall-Agent
Export-ModuleMember -Function Add-AppInstrumentation
Export-ModuleMember -Function Remove-AppInstrumentation
Export-ModuleMember -Function Test-AgentBitnessCoverage

<#
Export-ModuleMember -Function Test-AgentInstallation
Export-ModuleMember -Function Test-AgentRegistration
Export-ModuleMember -Function Test-AppInstrumentation
Export-ModuleMember -Function Test-GlobalInstrumentation
#>
