jh_o365_PIM.psm1

Set-StrictMode -Version Latest

$script:O365PimCachedRoles = @()

function Connect-O365Pim {
    [CmdletBinding()]
    param(
        [string[]]$Scopes = @(
            'User.Read',
            'Directory.Read.All',
            'RoleManagement.Read.Directory',
            'RoleEligibilitySchedule.Read.Directory',
            'RoleAssignmentSchedule.ReadWrite.Directory'
        ),
        [switch]$Force
    )

    if (-not (Get-Module -Name Microsoft.Graph.Authentication -ListAvailable)) {
        throw 'Microsoft Graph PowerShell SDK is not installed. Install it with: Install-Module Microsoft.Graph -Scope CurrentUser'
    }

    $context = $null
    try {
        $context = Get-MgContext -ErrorAction Stop
    }
    catch {
        $context = $null
    }

    if ($Force -or -not $context) {
        Write-Host 'Initiating login to Microsoft Graph for Entra PIM operations...' -ForegroundColor Cyan
        try {
            Connect-MgGraph -Scopes $Scopes -NoWelcome -ErrorAction Stop | Out-Null
            $context = Get-MgContext -ErrorAction Stop
        }
        catch {
            throw "Failed to connect to Microsoft Graph: $($_.Exception.Message)"
        }
    }
    else {
        $missingScopes = @($Scopes | Where-Object { $_ -notin $context.Scopes })
        if ($missingScopes.Count -gt 0) {
            Write-Host "Reconnecting to Microsoft Graph to add required scopes: $($missingScopes -join ', ')" -ForegroundColor Cyan
            try {
                Connect-MgGraph -Scopes $Scopes -NoWelcome -ErrorAction Stop | Out-Null
                $context = Get-MgContext -ErrorAction Stop
            }
            catch {
                throw "Failed to reconnect to Microsoft Graph: $($_.Exception.Message)"
            }
        }
    }

    if (-not $context) {
        throw 'Unable to establish a valid connection to Microsoft Graph. Please try again.'
    }

    Write-Host "Connected to Microsoft Graph as: $($context.Account)" -ForegroundColor Green
    return $context
}

function Get-O365PimEligibleRole {
    [CmdletBinding()]
    param(
        [switch]$PassThru
    )

    $context = Connect-O365Pim

    # Get user from context instead of using 'me' endpoint which may fail
    if (-not $context.Account) {
        throw 'Unable to determine the current signed-in user from the authentication context.'
    }

    $userAccount = $context.Account
    Write-Host "Fetching eligible roles for user: $userAccount" -ForegroundColor Cyan

    # Try to get the user object using the UPN from the context
    $me = $null
    try {
        $me = Get-MgUser -UserId $userAccount -ErrorAction Stop
    }
    catch {
        # If that fails, try using 'me'
        try {
            $me = Get-MgUser -UserId 'me' -ErrorAction Stop
        }
        catch {
            throw "Unable to retrieve current user information from Microsoft Graph. Verify you are logged in and have User.Read permission. Error: $($_.Exception.Message)"
        }
    }

    if (-not $me -or -not $me.Id) {
        throw 'Unable to determine the current signed-in user ID from Microsoft Graph.'
    }

    $schedules = Get-MgRoleManagementDirectoryRoleEligibilityScheduleInstance -All -Filter "principalId eq '$($me.Id)'"

    if (-not $schedules) {
        Write-Host 'No eligible Microsoft Entra roles were found for your account.' -ForegroundColor Yellow
        $script:O365PimCachedRoles = @()
        return @()
    }

    $roleMap = @{}
    $scopeMap = @{}

    foreach ($schedule in $schedules) {
        if (-not $roleMap.ContainsKey($schedule.RoleDefinitionId)) {
            $roleMap[$schedule.RoleDefinitionId] = Get-MgRoleManagementDirectoryRoleDefinition -UnifiedRoleDefinitionId $schedule.RoleDefinitionId
        }

        if (-not $scopeMap.ContainsKey($schedule.DirectoryScopeId)) {
            if ([string]::IsNullOrWhiteSpace($schedule.DirectoryScopeId) -or $schedule.DirectoryScopeId -eq '/') {
                $scopeMap[$schedule.DirectoryScopeId] = '/'
            }
            elseif ($schedule.DirectoryScopeId -like '/administrativeUnits/*') {
                $auId = $schedule.DirectoryScopeId.Split('/')[-1]
                try {
                    $au = Get-MgDirectoryAdministrativeUnit -AdministrativeUnitId $auId -ErrorAction Stop
                    $scopeMap[$schedule.DirectoryScopeId] = "Administrative Unit: $($au.DisplayName)"
                }
                catch {
                    $scopeMap[$schedule.DirectoryScopeId] = $schedule.DirectoryScopeId
                }
            }
            else {
                $scopeMap[$schedule.DirectoryScopeId] = $schedule.DirectoryScopeId
            }
        }
    }

    # Create temp objects with role names to enable alphabetical sorting
    $tempFormatted = @()
    foreach ($schedule in $schedules) {
        $roleName = $roleMap[$schedule.RoleDefinitionId].DisplayName
        if (-not $roleName) {
            continue
        }

        $scopeName = $scopeMap[$schedule.DirectoryScopeId]
        if (-not $scopeName) {
            $scopeName = '/'
        }

        $tempFormatted += [pscustomobject]@{
            RoleName         = $roleName
            Scope            = $scopeName
            RoleDefinitionId = $schedule.RoleDefinitionId
            DirectoryScopeId = $schedule.DirectoryScopeId
            AppScopeId       = $schedule.AppScopeId
            EligibilityId    = $schedule.Id
            PrincipalId      = $schedule.PrincipalId
        }
    }

    # Sort alphabetically by role name and assign numbers
    $formatted = @()
    $i = 1

    foreach ($item in ($tempFormatted | Sort-Object RoleName)) {
        $formatted += [pscustomobject]@{
            Number           = '{0:d2}' -f $i
            RoleName         = $item.RoleName
            Scope            = $item.Scope
            RoleDefinitionId = $item.RoleDefinitionId
            DirectoryScopeId = $item.DirectoryScopeId
            AppScopeId       = $item.AppScopeId
            EligibilityId    = $item.EligibilityId
            PrincipalId      = $item.PrincipalId
        }

        $i++
    }

    $script:O365PimCachedRoles = $formatted

    Write-Host "Found $($formatted.Count) eligible role(s)" -ForegroundColor Green
    Write-Host ""
    
    # Display the table
    $formatted | Select-Object Number, RoleName, Scope | Format-Table -AutoSize | Out-String | Write-Host

    if ($PassThru) {
        return , $formatted
    }
}

function Get-O365PimActiveAssignment {
    [CmdletBinding()]
    param(
        [switch]$PassThru
    )

    $context = Connect-O365Pim

    if (-not $context.Account) {
        throw 'Unable to determine the current signed-in user from the authentication context.'
    }

    $userAccount = $context.Account
    Write-Host "Fetching active role assignments for user: $userAccount" -ForegroundColor Cyan

    $me = $null
    try {
        $me = Get-MgUser -UserId $userAccount -ErrorAction Stop
    }
    catch {
        try {
            $me = Get-MgUser -UserId 'me' -ErrorAction Stop
        }
        catch {
            throw "Unable to retrieve current user information from Microsoft Graph. Error: $($_.Exception.Message)"
        }
    }

    if (-not $me -or -not $me.Id) {
        throw 'Unable to determine the current signed-in user ID from Microsoft Graph.'
    }

    # Get active role assignment schedules (not eligible, but actual active assignments)
    $assignments = Get-MgRoleManagementDirectoryRoleAssignmentScheduleInstance -All -Filter "principalId eq '$($me.Id)'"

    if (-not $assignments) {
        Write-Host 'No active role assignments found.' -ForegroundColor Yellow
        return @()
    }

    Write-Host "Found $($assignments.Count) active assignment(s)" -ForegroundColor Green
    Write-Host ""

    $roleMap = @{}
    $tempFormatted = @()

    foreach ($assignment in $assignments) {
        if (-not $roleMap.ContainsKey($assignment.RoleDefinitionId)) {
            $roleMap[$assignment.RoleDefinitionId] = Get-MgRoleManagementDirectoryRoleDefinition -UnifiedRoleDefinitionId $assignment.RoleDefinitionId
        }

        $roleName = $roleMap[$assignment.RoleDefinitionId].DisplayName
        $scope = if ([string]::IsNullOrWhiteSpace($assignment.DirectoryScopeId) -or $assignment.DirectoryScopeId -eq '/') { '/' } else { $assignment.DirectoryScopeId }

        $tempFormatted += [pscustomobject]@{
            RoleName         = $roleName
            Scope            = $scope
            StartDateTime    = $assignment.StartDateTime
            EndDateTime      = $assignment.EndDateTime
            RoleDefinitionId = $assignment.RoleDefinitionId
            DirectoryScopeId = $assignment.DirectoryScopeId
        }
    }

    $formatted = @()
    $i = 1
    foreach ($item in ($tempFormatted | Sort-Object RoleName)) {
        $formatted += [pscustomobject]@{
            Number           = '{0:d2}' -f $i
            RoleName         = $item.RoleName
            Scope            = $item.Scope
            StartDateTime    = $item.StartDateTime
            EndDateTime      = $item.EndDateTime
            RoleDefinitionId = $item.RoleDefinitionId
            DirectoryScopeId = $item.DirectoryScopeId
        }
        $i++
    }

    $formatted | Select-Object Number, RoleName, Scope, StartDateTime, EndDateTime | Format-Table -AutoSize | Out-String | Write-Host

    if ($PassThru) {
        return , $formatted
    }
}


function Enable-O365PimRole {
    [CmdletBinding(DefaultParameterSetName = 'ByNumber', SupportsShouldProcess)]
    param(
        [Parameter(Mandatory, ParameterSetName = 'ByNumber')]
        [ValidatePattern('^\d+$')]
        [string[]]$RoleNumber,

        [Parameter(Mandatory, ParameterSetName = 'ByName')]
        [string[]]$RoleName,

        [string]$Reason = 'STB',

        [ValidateRange(1, 24)]
        [int]$DurationHours = 8,

        [switch]$RefreshRoleList
    )

    Connect-O365Pim | Out-Null

    $roles = $script:O365PimCachedRoles
    if ($RefreshRoleList -or -not $roles -or $roles.Count -eq 0) {
        $roles = Get-O365PimEligibleRole -PassThru
    }

    if (-not $roles -or $roles.Count -eq 0) {
        throw 'No eligible roles are available to activate for the signed-in user.'
    }

    $rolesToActivate = @()

    if ($PSCmdlet.ParameterSetName -eq 'ByNumber') {
        foreach ($num in $RoleNumber) {
            $normalizedNumber = '{0:d2}' -f [int]$num
            $targetRole = $roles | Where-Object { $_.Number -eq $normalizedNumber }
            if (-not $targetRole) {
                Write-Host "Warning: Role number '$num' was not found. Skipping." -ForegroundColor Yellow
                continue
            }
            $rolesToActivate += $targetRole
        }
    }
    else {
        foreach ($name in $RoleName) {
            $matches = @($roles | Where-Object { $_.RoleName -like $name })
            if ($matches.Count -eq 0) {
                $matches = @($roles | Where-Object { $_.RoleName -eq $name })
            }

            if ($matches.Count -eq 0) {
                Write-Host "Warning: Role name '$name' was not found. Skipping." -ForegroundColor Yellow
                continue
            }

            if ($matches.Count -gt 1) {
                Write-Host "Warning: Role name '$name' matched multiple entries. Use -RoleNumber instead. Skipping." -ForegroundColor Yellow
                continue
            }

            $rolesToActivate += $matches[0]
        }
    }

    if ($rolesToActivate.Count -eq 0) {
        throw 'No valid roles were found to activate.'
    }

    Write-Host "Activating $($rolesToActivate.Count) role(s)..." -ForegroundColor Cyan
    Write-Host ""

    $results = @()

    foreach ($targetRole in $rolesToActivate) {
        $activationTarget = "$($targetRole.Number) - $($targetRole.RoleName) ($($targetRole.Scope))"
        if (-not $PSCmdlet.ShouldProcess($activationTarget, 'Activate PIM role')) {
            continue
        }

        $now = (Get-Date).ToUniversalTime().ToString('o')
        $duration = "PT${DurationHours}H"

        $body = @{
            Action           = 'selfActivate'
            PrincipalId      = $targetRole.PrincipalId
            RoleDefinitionId = $targetRole.RoleDefinitionId
            DirectoryScopeId = if ([string]::IsNullOrWhiteSpace($targetRole.DirectoryScopeId)) { '/' } else { $targetRole.DirectoryScopeId }
            Justification    = $Reason
            ScheduleInfo     = @{
                StartDateTime = $now
                Expiration    = @{
                    Type     = 'AfterDuration'
                    Duration = $duration
                }
            }
        }

        if (-not [string]::IsNullOrWhiteSpace($targetRole.AppScopeId)) {
            $body.AppScopeId = $targetRole.AppScopeId
        }

        try {
            $request = New-MgRoleManagementDirectoryRoleAssignmentScheduleRequest -BodyParameter $body -ErrorAction Stop
            
            $result = [pscustomobject]@{
                RoleNumber      = $targetRole.Number
                RoleName        = $targetRole.RoleName
                Scope           = $targetRole.Scope
                Reason          = $Reason
                DurationHours   = $DurationHours
                RequestId       = $request.Id
                RequestStatus   = $request.Status
                RequestedAtUtc  = $request.CreatedDateTime
            }

            $results += $result

            Write-Host "✓ Activation request submitted for role [$($targetRole.Number)] $($targetRole.RoleName)" -ForegroundColor Green
        }
        catch {
            $errorMsg = $_.Exception.Message
            
            if ($errorMsg -like '*RoleAssignmentExists*') {
                Write-Host "✗ Role [$($targetRole.Number)] $($targetRole.RoleName) already has an active or pending assignment." -ForegroundColor Yellow
                Write-Host " → Check Get-O365PimActiveAssignment to see current active assignments" -ForegroundColor Gray
                Write-Host " → If not shown there, the assignment may have expired but the schedule entry still exists" -ForegroundColor Gray
            }
            elseif ($errorMsg -like '*PendingRoleAssignmentRequest*') {
                Write-Host "✗ Role [$($targetRole.Number)] $($targetRole.RoleName) has a pending activation request." -ForegroundColor Yellow
                Write-Host " → Wait for approval or the request to complete before activating again" -ForegroundColor Gray
            }
            elseif ($errorMsg -like '*RoleAssignmentRequestPolicyValidationFailed*') {
                Write-Host "✗ Role [$($targetRole.Number)] $($targetRole.RoleName) activation failed due to policy validation." -ForegroundColor Red
                Write-Host " → Error: $errorMsg" -ForegroundColor Gray
                Write-Host " → Try a shorter duration (e.g., -DurationHours 4)" -ForegroundColor Gray
            }
            else {
                Write-Host "✗ Failed to activate role [$($targetRole.Number)] $($targetRole.RoleName): $errorMsg" -ForegroundColor Red
            }
        }
    }

    Write-Host ""
    Write-Host "Summary: Activated $($results.Count) of $($rolesToActivate.Count) role(s)" -ForegroundColor Cyan
    Write-Host "Reason: $Reason" -ForegroundColor Cyan

    return $results
}

function Disable-O365PimRole {
    [CmdletBinding(DefaultParameterSetName = 'ByNumber', SupportsShouldProcess)]
    param(
        [Parameter(Mandatory, ParameterSetName = 'ByNumber')]
        [ValidatePattern('^\d+$')]
        [string[]]$RoleNumber,

        [Parameter(Mandatory, ParameterSetName = 'ByName')]
        [string[]]$RoleName,

        [switch]$RefreshRoleList
    )

    Connect-O365Pim | Out-Null

    $roles = $script:O365PimCachedRoles
    if ($RefreshRoleList -or -not $roles -or $roles.Count -eq 0) {
        $roles = Get-O365PimEligibleRole -PassThru
    }

    if (-not $roles -or $roles.Count -eq 0) {
        throw 'No eligible roles are available to deactivate for the signed-in user.'
    }

    $rolesToDeactivate = @()

    if ($PSCmdlet.ParameterSetName -eq 'ByNumber') {
        foreach ($num in $RoleNumber) {
            $normalizedNumber = '{0:d2}' -f [int]$num
            $targetRole = $roles | Where-Object { $_.Number -eq $normalizedNumber }
            if (-not $targetRole) {
                Write-Host "Warning: Role number '$num' was not found. Skipping." -ForegroundColor Yellow
                continue
            }
            $rolesToDeactivate += $targetRole
        }
    }
    else {
        foreach ($name in $RoleName) {
            $matches = @($roles | Where-Object { $_.RoleName -like $name })
            if ($matches.Count -eq 0) {
                $matches = @($roles | Where-Object { $_.RoleName -eq $name })
            }

            if ($matches.Count -eq 0) {
                Write-Host "Warning: Role name '$name' was not found. Skipping." -ForegroundColor Yellow
                continue
            }

            if ($matches.Count -gt 1) {
                Write-Host "Warning: Role name '$name' matched multiple entries. Use -RoleNumber instead. Skipping." -ForegroundColor Yellow
                continue
            }

            $rolesToDeactivate += $matches[0]
        }
    }

    if ($rolesToDeactivate.Count -eq 0) {
        throw 'No valid roles were found to deactivate.'
    }

    Write-Host "Deactivating $($rolesToDeactivate.Count) role(s)..." -ForegroundColor Cyan
    Write-Host ""

    $results = @()

    foreach ($targetRole in $rolesToDeactivate) {
        $deactivationTarget = "$($targetRole.Number) - $($targetRole.RoleName) ($($targetRole.Scope))"
        if (-not $PSCmdlet.ShouldProcess($deactivationTarget, 'Deactivate PIM role')) {
            continue
        }

        try {
            # Deactivate by creating a request with 'SelfDeactivate' action
            $body = @{
                Action           = 'SelfDeactivate'
                PrincipalId      = $targetRole.PrincipalId
                RoleDefinitionId = $targetRole.RoleDefinitionId
                DirectoryScopeId = if ([string]::IsNullOrWhiteSpace($targetRole.DirectoryScopeId)) { '/' } else { $targetRole.DirectoryScopeId }
            }

            $request = New-MgRoleManagementDirectoryRoleAssignmentScheduleRequest -BodyParameter $body -ErrorAction Stop

            $result = [pscustomobject]@{
                RoleNumber      = $targetRole.Number
                RoleName        = $targetRole.RoleName
                Scope           = $targetRole.Scope
                RequestId       = $request.Id
                RequestStatus   = $request.Status
                RequestedAtUtc  = $request.CreatedDateTime
            }

            $results += $result

            Write-Host "✓ Deactivation request submitted for role [$($targetRole.Number)] $($targetRole.RoleName)" -ForegroundColor Green
        }
        catch {
            $errorMsg = $_.Exception.Message
            
            if ($errorMsg -like '*RoleAssignmentNotFound*' -or $errorMsg -like '*does not exist*') {
                Write-Host "✗ Role [$($targetRole.Number)] $($targetRole.RoleName) is not currently active." -ForegroundColor Yellow
            }
            else {
                Write-Host "✗ Failed to deactivate role [$($targetRole.Number)] $($targetRole.RoleName): $errorMsg" -ForegroundColor Red
            }
        }
    }

    Write-Host ""
    Write-Host "Summary: Deactivation requested for $($results.Count) of $($rolesToDeactivate.Count) role(s)" -ForegroundColor Cyan

    return $results
}

Export-ModuleMember -Function Connect-O365Pim, Get-O365PimEligibleRole, Get-O365PimActiveAssignment, Enable-O365PimRole, Disable-O365PimRole