<#
.SYNOPSIS
  Find machines by tag in Microsoft Defender for Endpoint and offboard them via API.

.NOTES
  - Requires an Azure AD app (client id + secret) with Application permissions:
      - Machine.Read.All (or Machine.ReadWrite.All) to list machines by tag
      - Machine.Offboard to offboard machines
    Grant admin consent for the permissions.
  - Choose the correct API base for your geography (default below is eu). Example hosts:
      - us.api.security.microsoft.com
      - eu.api.security.microsoft.com
      - uk.api.security.microsoft.com
      - api.securitycenter.microsoft.com (other docs reference this variant)
    See MS docs for recommended regional endpoint for best performance.
  - The Offboard API requires a non-empty "Comment" JSON body.
  - Script runs in dry-run mode by default; add -Offboard to actually submit offboard requests.
#>

param(
    [Parameter()] [string] $TenantId,
    [Parameter()] [string] $ClientId,
    [Parameter()] [string] $ClientSecret,
    [Parameter(HelpMessage="The machine tag to find (exact match unless -StartsWith used)")] [string] $Tag,
    [Parameter()] [string] $CompletedTag = "offboarded",
    [Parameter()] [switch] $StartsWith,
    [Parameter()] [string] $ApiBase = "https://eu.api.security.microsoft.com", # change to your region
    [Parameter()] [switch] $Offboard,  # only offboard when explicitly requested
    [Parameter()] [int] $DelayBetweenCallsSec = 1,   # breathing room for rate limits
    [Parameter()] [int] $MaxRetries = 3,
    [Parameter()] [int] $RetryDelaySec = 5,
    [Parameter(HelpMessage="Optional path to a JSON or PSD1 file that stores parameter values. CLI args override file entries.")] [string] $ParametersPath
)

function ConvertTo-Hashtable {
    param(
        [Parameter()] [psobject] $InputObject
    )

    if ($null -eq $InputObject) {
        return @{}
    }

    if ($InputObject -is [hashtable]) {
        return $InputObject
    }

    $hash = @{}
    foreach ($prop in $InputObject.PSObject.Properties) {
        $hash[$prop.Name] = $prop.Value
    }
    return $hash
}

function Import-ParameterFile {
    param(
        [Parameter(Mandatory=$true)] [string] $Path
    )

    if (-not (Test-Path -LiteralPath $Path)) {
        throw "Parameter file '$Path' not found."
    }

    $resolvedPath = (Resolve-Path -LiteralPath $Path).ProviderPath
    $extension = [System.IO.Path]::GetExtension($resolvedPath).ToLowerInvariant()

    try {
        switch ($extension) {
            ".psd1" {
                $data = Import-PowerShellDataFile -Path $resolvedPath
            }
            default {
                $raw = Get-Content -LiteralPath $resolvedPath -Raw
                if ([string]::IsNullOrWhiteSpace($raw)) {
                    return @{}
                }
                $json = ConvertFrom-Json -InputObject $raw
                $data = ConvertTo-Hashtable -InputObject $json
            }
        }

        return ConvertTo-Hashtable -InputObject $data
    } catch {
        throw "Failed to parse parameter file '$resolvedPath': $($_.Exception.Message)"
    }
}

function Get-ScriptRoot {
    if ($null -ne $PSScriptRoot) {
        if ($PSScriptRoot -is [string] -and -not [string]::IsNullOrWhiteSpace($PSScriptRoot)) {
            return [string]$PSScriptRoot
        }

        if ($PSScriptRoot -is [System.Collections.IEnumerable]) {
            foreach ($item in $PSScriptRoot) {
                if (-not [string]::IsNullOrWhiteSpace([string]$item)) {
                    return [string]$item
                }
            }
        }
    }

    if ($MyInvocation -and $MyInvocation.MyCommand -and $MyInvocation.MyCommand.Path) {
        return Split-Path -Parent $MyInvocation.MyCommand.Path
    }

    return (Get-Location).ProviderPath
}

function Get-AccessToken {
    param(
        [string] $TenantId,
        [string] $ClientId,
        [string] $ClientSecret,
        [string] $Scope = "https://api.securitycenter.microsoft.com/.default"
    )
    $tokenEndpoint = "https://login.microsoftonline.com/$TenantId/oauth2/v2.0/token"

    $body = @{
        client_id     = $ClientId
        client_secret = $ClientSecret
        scope         = $Scope
        grant_type    = "client_credentials"
    }

    try {
        $resp = Invoke-RestMethod -Method Post -Uri $tokenEndpoint -Body $body -ErrorAction Stop
        return $resp.access_token
    } catch {
        Write-Error "Failed to obtain access token: $($_.Exception.Message)"
        throw
    }
}

function Find-MachinesByTag {
    param(
        [string] $ApiBase,
        [string] $Token,
        [string] $Tag,
        [bool] $StartsWith = $false
    )

    $uri = "$($ApiBase.TrimEnd('/'))/api/machines/findbytag?tag=$([uri]::EscapeDataString($Tag))&useStartsWithFilter=$($StartsWith.ToString().ToLower())"

    $headers = @{
        Authorization = "Bearer $Token"
    }

    try {
        $resp = Invoke-RestMethod -Method Get -Uri $uri -Headers $headers -ErrorAction Stop
        if ($null -eq $resp) {
            return @()
        }

        # Some Defender APIs wrap results under a 'value' property; others return a bare array.
        $machines = if ($resp.PSObject.Properties.Match('value').Count -gt 0) {
            $resp.value
        } else {
            $resp
        }

        if ($null -eq $machines) {
            return @()
        }

        if ($machines -is [string] -or -not ($machines -is [System.Collections.IEnumerable])) {
            $machines = @($machines)
        }

        return @($machines)
    } catch {
        Write-Error "Find-by-tag call failed: $($_.Exception.Message)"
        throw
    }
}

function Invoke-MachineOffboard {
    param(
        [string] $ApiBase,
        [string] $Token,
        [string] $MachineId,
        [string] $Comment = "Offboarded by automation",
        [int] $MaxRetries = 3,
        [int] $RetryDelaySec = 5
    )

    $uri = "$($ApiBase.TrimEnd('/'))/api/machines/$MachineId/offboard"
    $headers = @{
        Authorization = "Bearer $Token"
        "Content-Type" = "application/json"
    }
    $body = @{ Comment = $Comment } | ConvertTo-Json

    for ($attempt = 1; $attempt -le $MaxRetries; $attempt++) {
        try {
            $resp = Invoke-RestMethod -Method Post -Uri $uri -Headers $headers -Body $body -ErrorAction Stop
            return @{ Success = $true; Response = $resp }
        } catch {
            $status = if ($_.Exception.Response) {
                try { ($_.Exception.Response | Select-Object -ExpandProperty StatusCode).Value__ } catch { $_.Exception.Message }
            } else { $_.Exception.Message }
            Write-Warning "Offboard attempt $attempt failed for $MachineId (status: $status)."
            if ($attempt -lt $MaxRetries) {
                Start-Sleep -Seconds $RetryDelaySec
            } else {
                return @{ Success = $false; Error = $_.Exception.Message }
            }
        }
    }
}

function Update-MachineTag {
    param(
        [string] $ApiBase,
        [string] $Token,
        [string] $MachineId,
        [string] $TagValue,
        [ValidateSet("Add","Remove")] [string] $Action
    )

    if ([string]::IsNullOrWhiteSpace($TagValue)) {
        return
    }

    $uri = "$($ApiBase.TrimEnd('/'))/api/machines/$MachineId/tags"
    $headers = @{
        Authorization = "Bearer $Token"
        "Content-Type" = "application/json"
    }
    $body = @{
        Value  = $TagValue
        Action = $Action
    } | ConvertTo-Json

    Invoke-RestMethod -Method Post -Uri $uri -Headers $headers -Body $body -ErrorAction Stop | Out-Null
}

function Set-OffboardCompletionTag {
    param(
        [string] $ApiBase,
        [string] $Token,
        [string] $MachineId,
        [string] $OriginalTag,
        [string] $CompletedTag
    )

    if ([string]::IsNullOrWhiteSpace($CompletedTag)) {
        return
    }

    if (-not [string]::IsNullOrWhiteSpace($OriginalTag) -and $OriginalTag -ne $CompletedTag) {
        try {
            Update-MachineTag -ApiBase $ApiBase -Token $Token -MachineId $MachineId -TagValue $OriginalTag -Action "Remove"
            Write-Host "Removed tag '$OriginalTag' from $MachineId."
        } catch {
        Write-Warning "Failed to remove tag '$OriginalTag' from ${MachineId}: $($_.Exception.Message)"
        }
    }

    try {
        Update-MachineTag -ApiBase $ApiBase -Token $Token -MachineId $MachineId -TagValue $CompletedTag -Action "Add"
        Write-Host "Added tag '$CompletedTag' to $MachineId."
    } catch {
        Write-Warning "Failed to add tag '$CompletedTag' to ${MachineId}: $($_.Exception.Message)"
    }
}

function Get-ActiveOffboardRequest {
    param(
        [string] $ApiBase,
        [string] $Token,
        [string] $MachineId
    )

    $filter = "type eq 'Offboard' and machineId eq '$MachineId'"
    $query = "`$filter=$([uri]::EscapeDataString($filter))&`$orderby=creationDateTime%20desc&`$top=5"
    $uri = "$($ApiBase.TrimEnd('/'))/api/machineactions?$query"

    $headers = @{
        Authorization = "Bearer $Token"
    }

    try {
        $resp = Invoke-RestMethod -Method Get -Uri $uri -Headers $headers -ErrorAction Stop
    } catch {
        Write-Warning "Failed to query machine actions for ${MachineId}: $($_.Exception.Message)"
        return $null
    }

    if ($null -eq $resp) {
        return $null
    }

    $actions = if ($resp.PSObject.Properties.Match('value').Count -gt 0) {
        $resp.value
    } else {
        $resp
    }

    if ($null -eq $actions) {
        return $null
    }

    if ($actions -isnot [System.Collections.IEnumerable]) {
        $actions = @($actions)
    }

    $completedStatuses = @('Succeeded','Failed','Cancelled','TimedOut')
    foreach ($action in $actions) {
        $status = $action.status
        if ([string]::IsNullOrWhiteSpace($status)) {
            return $action
        }

        if ($completedStatuses -notcontains $status) {
            return $action
        }
    }

    return $null
}

# Allow parameters to be stored in a reusable file instead of always passing them via CLI.
$scriptRoot = Get-ScriptRoot
$parameterDefaults = @{}
if ([string]::IsNullOrWhiteSpace($ParametersPath)) {
    $defaultCandidates = @(
        (Join-Path -Path $scriptRoot -ChildPath "MDE_OffboardDevices.parameters.json")
        (Join-Path -Path $scriptRoot -ChildPath "MDE_OffboardDevices.parameters.psd1")
    )
    foreach ($candidate in $defaultCandidates) {
        if (Test-Path -LiteralPath $candidate) {
            $ParametersPath = $candidate
            break
        }
    }
}

if (-not [string]::IsNullOrWhiteSpace($ParametersPath)) {
    Write-Host "Loading parameter defaults from '$ParametersPath'. CLI values take precedence."
    $parameterDefaults = Import-ParameterFile -Path $ParametersPath
}

if ($parameterDefaults.Count -gt 0) {
    if ($parameterDefaults.ContainsKey('WhatIf') -and -not $parameterDefaults.ContainsKey('Offboard')) {
        Write-Verbose "Detected legacy 'WhatIf' setting in parameter file. Translating to new -Offboard logic."
        $parameterDefaults['Offboard'] = -not [bool]$parameterDefaults['WhatIf']
    }

    $stringParams = @('TenantId','ClientId','ClientSecret','Tag','CompletedTag','ApiBase')
    foreach ($name in $stringParams) {
        if (-not $PSBoundParameters.ContainsKey($name) -and $parameterDefaults.ContainsKey($name)) {
            Set-Variable -Name $name -Value ([string]$parameterDefaults[$name]) -Scope Script
        }
    }

    $intParams = @('DelayBetweenCallsSec','MaxRetries','RetryDelaySec')
    foreach ($name in $intParams) {
        if (-not $PSBoundParameters.ContainsKey($name) -and $parameterDefaults.ContainsKey($name) -and $null -ne $parameterDefaults[$name]) {
            Set-Variable -Name $name -Value ([int]$parameterDefaults[$name]) -Scope Script
        }
    }

    $switchParams = @('StartsWith','Offboard')
    foreach ($name in $switchParams) {
        if (-not $PSBoundParameters.ContainsKey($name) -and $parameterDefaults.ContainsKey($name)) {
            Set-Variable -Name $name -Value ([bool]$parameterDefaults[$name]) -Scope Script
        }
    }
}

$missingRequired = @()
foreach ($req in @('TenantId','ClientId','ClientSecret','Tag')) {
    $value = Get-Variable -Name $req -ValueOnly
    if ([string]::IsNullOrWhiteSpace([string]$value)) {
        $missingRequired += $req
    }
}

if ($missingRequired.Count -gt 0) {
    throw "Missing required parameter(s): $($missingRequired -join ', '). Provide them via CLI arguments or a parameter file."
}

# main
Write-Host "Using API base: $ApiBase"
$token = Get-AccessToken -TenantId $TenantId -ClientId $ClientId -ClientSecret $ClientSecret

Write-Host "Searching for machines with tag '$Tag' (StartsWith = $StartsWith)..."
$machines = Find-MachinesByTag -ApiBase $ApiBase -Token $token -Tag $Tag -StartsWith $StartsWith

if (-not $machines -or $machines.Count -eq 0) {
    Write-Host "No machines found for tag '$Tag'. Exiting."
    return
}

Write-Host "Found $($machines.Count) machines. Listing IDs..."
$machines | ForEach-Object {
    $id = $_.id
    $dns = $_.computerDnsName
    $deviceTags = if ($_.machineTags) { ($_.machineTags -join ', ') } else { 'none' }
    Write-Host " - ID: $id  DNS: $dns  Tags: $deviceTags"
}

if (-not $Offboard) {
    Write-Host "`nDry run: -Offboard was not supplied, so no offboard calls were sent."
    Write-Host "Re-run with -Offboard to offboard the listed machines."
    return
}

foreach ($m in $machines) {
    $id = $m.id
    if (-not $id) {
        Write-Warning "Machine entry missing 'id' property, skipping: $($m | ConvertTo-Json -Depth 3)"
        continue
    }

    $activeRequest = Get-ActiveOffboardRequest -ApiBase $ApiBase -Token $token -MachineId $id
    if ($null -ne $activeRequest) {
        $status = $activeRequest.status
        $requestedOn = $activeRequest.creationDateTime
        Write-Warning "Offboard already requested for $id (status: $status, requested: $requestedOn). Skipping."
        continue
    }

    Write-Host "`nOffboarding machine $id ..."
    $result = Invoke-MachineOffboard -ApiBase $ApiBase -Token $token -MachineId $id -Comment "Offboard by automation (tag: $Tag)" -MaxRetries $MaxRetries -RetryDelaySec $RetryDelaySec

    if ($result.Success) {
        Write-Host "Offboard requested for $id. Response: "
        $result.Response | ConvertTo-Json -Depth 5 | Write-Host

        if (-not [string]::IsNullOrWhiteSpace($CompletedTag)) {
            Set-OffboardCompletionTag -ApiBase $ApiBase -Token $token -MachineId $id -OriginalTag $Tag -CompletedTag $CompletedTag
        }

    } else {
        Write-Error "Failed to offboard $id. Error: $($result.Error)"
    }

    Start-Sleep -Seconds $DelayBetweenCallsSec
}

Write-Host "`nDone."