#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"
    $dll = $dlPath + "/AppsentinelsIISModule.dll"

    Prepare-Install-Directory -InstallDir $installDir

    try {
        Copy-Item $dll -Destination $installDir
        [System.Environment]::SetEnvironmentVariable('APPSENTINELS_AGENT_INSTALL_DIR', $installDir, [System.EnvironmentVariableTarget]::Machine)
        Write-Output "Installed Appsentinels IIS Agent Successfully"
    }
    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 {
	return (Module-Exists "AppsentinelsIISAgent")
}

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
}

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

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

    Add-WebConfiguration -Filter "/$configPath" -Value @{name='AppsentinelsIISAgent'; image=$dll; preCondition='bitness64'}

    Write-Output "Appsentinels IIS Agent added as IIS native module"
}


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

    Reset-IIS
}

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."
    }

    # nothing to do return. All the websites would inherit global modules
    # and Check-If-Website-Instrumentation-Exists would return true even in case of global instrumentation
<#
    if (Test-GlobalInstrumentation) {
        Write-Output "Global Instrumentation already exists."
        return
    }


    $result = Check-If-Website-Instrumentation-Exists
    if ($result.IsInstrumented) {
        if (-not $Force) {
            Write-Error "App Instrumentation already exists. Remove first or use -Force"
        } else {
            if ($result.IsInstrumented) {
                foreach ($site in $result.Sites) {
                    Remove-AppInstrumentation -website $site
                }
            }
        }
    }
#>
    $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."
        }
    }
    Update-WebModuleConfiguration -ModuleName "AppsentinelsIISAgent" -PreCondition "bitness64"
    Write-Output "IIS is instrumented globally"
}

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

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

    $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
    $dllPath = Join-Path $installDir "AppsentinelsIISModule.dll"

    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 "App 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."
    }

    & $appCmd --% clear config -section:system.webServer/globalModules /[name='AppsentinelsIISAgent',preCondition='bitness64'] /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() {
    if(Test-GlobalInstrumentation) {
        Remove-GlobalInstrumentation
    }

<#
    $result = Check-If-Website-Instrumentation-Exists
    if ($result.IsInstrumented) {
        foreach ($site in $result.Sites) {
            Remove-AppInstrumentation -website $site
        }
    }
#>
}

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
    }

    Clear-WebConfiguration -PSPath 'MACHINE/WEBROOT/APPHOST' -Filter "//modules/add[@name='AppsentinelsIISAgent']"

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

function Test-GlobalInstrumentation {
    $configPath = "system.webServer/modules"
    $module = Get-WebConfiguration -PSPath "MACHINE/WEBROOT/APPHOST" -Filter "/$configPath/add[@name='AppsentinelsIISAgent']"
    return ($module -ne $null)
}

function Add-AppInstrumentation() {
    param(
        [Parameter(Mandatory = $true)]
        [string]$website,
        [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
    }

    $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."
        }
    }
    Update-WebModuleConfiguration -ModuleName "AppsentinelsIISAgent" -PreCondition "bitness64" -Website $website
    Write-Output "IIS is instrumented for website $website"
}

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

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

    Clear-WebConfiguration -PSPath "IIS:\Sites\$Website" -Filter "//modules/add[@name='AppsentinelsIISAgent']"
    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-AgentInstallation
Export-ModuleMember -Function Test-AgentRegistration
Export-ModuleMember -Function Test-AppInstrumentation
Export-ModuleMember -Function Test-GlobalInstrumentation
#>
