source/Private/ServicePrincipalsOperations.ps1

function Get-ServicePrincipalByAppId {
    param (
        [Parameter(Mandatory = $true)]
        [string]$AppId
    )

    Write-Verbose "Retrieving service principal for AppId $AppId"

    $servicePrincipal = Invoke-RJMgGraphCommand { Get-MgServicePrincipalByAppId -AppId $AppId -Property Id,DisplayName,AppRoles }

    if ($servicePrincipal) {
        Write-Verbose "Service principal exists: $($servicePrincipal.DisplayName) (ObjectId: $($servicePrincipal.Id))"
        return $servicePrincipal
    }

    Write-Verbose "No existing service principal found for AppId $AppId"
    return $null
}

function Get-CurrentAppRoleAssignmentList {
    param (
        [Parameter(Mandatory = $true)]
        [string]$ServicePrincipalId
    )

    Write-Verbose "Retrieving current app role assignments for service principal $ServicePrincipalId"

    $assignments = Invoke-RJMgGraphCommand { Get-MgServicePrincipalAppRoleAssignment -ServicePrincipalId $ServicePrincipalId }

    if ($assignments) {
        # Ensure we always return an array, even for single results
        # Use the comma operator to force array output
        return ,@($assignments)
    }

    return @()
}

function Remove-AppRoleAssignment {
    [CmdletBinding(SupportsShouldProcess)]
    param (
        [Parameter(Mandatory = $true)]
        [PSCustomObject]$Assignment
    )

    if ($PSCmdlet.ShouldProcess("App Role Assignment $($Assignment.AppRoleId) from Resource $($Assignment.ResourceId)", "Remove")) {
        Write-Verbose "Removing app role assignment: $($Assignment.AppRoleId) from resource $($Assignment.ResourceId)"

        Invoke-RJMgGraphCommand { Remove-MgServicePrincipalAppRoleAssignedTo -ServicePrincipalId $Assignment.ResourceId -AppRoleAssignmentId $Assignment.Id} | Out-Null

        Write-Information "Removed app role assignment: $($Assignment.AppRoleId)"
    }
}

function Add-AppRoleAssignment {
    param (
        [Parameter(Mandatory = $true)]
        [string]$ServicePrincipalId,
        [Parameter(Mandatory = $true)]
        [string]$ResourceId,
        [Parameter(Mandatory = $true)]
        [string]$AppRoleId,
        [Parameter(Mandatory = $true)]
        [string]$PermissionName
    )

    Write-Verbose "Adding app role assignment: $AppRoleId to service principal $ServicePrincipalId"

    $assignmentParams = @{
        ServicePrincipalId = $ResourceId
        BodyParameter = @{
            principalId = $ServicePrincipalId
            resourceId = $ResourceId
            appRoleId = $AppRoleId
        }
    }

    Invoke-RJMgGraphCommand { New-MgServicePrincipalAppRoleAssignedTo @assignmentParams} | Out-Null

    Write-Information "Added app role assignment: $AppRoleId ($PermissionName)"
}

function Invoke-AppRoleAssignmentChange {

    param (
        [Parameter(Mandatory = $true)]
        [string]$ServicePrincipalId,
        [hashtable[]]$AssignmentsToAdd = @(),
        [PSCustomObject[]]$AssignmentsToRemove = @()
    )

    # Remove unwanted assignments first
    if ($AssignmentsToRemove.Count -gt 0) {
        Write-Verbose "Removing $($AssignmentsToRemove.Count) unwanted assignments"

        foreach ($assignmentToRemove in $AssignmentsToRemove) {
            Remove-AppRoleAssignment -Assignment $assignmentToRemove
        }
    }
    else {
        Write-Verbose "No assignments to remove"
    }

    # Add new assignments
    if ($AssignmentsToAdd.Count -gt 0) {
        Write-Verbose "Adding $($AssignmentsToAdd.Count) new assignments"

        foreach ($assignmentToAdd in $AssignmentsToAdd) {
            $resourceId = $assignmentToAdd.ResourceId
            $appRoleId = $assignmentToAdd.AppRoleId

            Write-Verbose "Adding assignment: resourceId=$resourceId, appRoleId=$appRoleId"
            Add-AppRoleAssignment -ServicePrincipalId $ServicePrincipalId -ResourceId $resourceId -AppRoleId $appRoleId -PermissionName $assignmentToAdd.PermissionName
        }
    }
    else {
        Write-Verbose "No assignments to add"
    }
}

function Set-RJServicePrincipal {
    [CmdletBinding(SupportsShouldProcess)]
    param(
        [Parameter(Mandatory = $true)]
        [string]$AppId,
        [hashtable[]]$AssignmentsToAdd = @(),
        [PSCustomObject[]]$AssignmentsToRemove = @()
    )

    if ($PSCmdlet.ShouldProcess("Service Principal for App ID $AppId", "Create/Update with permissions")) {
        # Get or create service principal
        Write-Verbose "Checking if service principal already exists for AppId $AppId"
        $servicePrincipal = Get-ServicePrincipalByAppId -AppId $AppId

        if (-not $servicePrincipal) {
            Write-Verbose "Creating service principal for AppId $AppId"

            $servicePrincipal = Invoke-RJMgGraphCommand { New-MgServicePrincipal -AppId $AppId -Tags @('WindowsAzureActiveDirectoryIntegratedApp') }

            Write-Information "Created service principal: $($servicePrincipal.DisplayName) (ObjectId: $($servicePrincipal.Id))"
        } else {
            Write-Information "Service principal exists: $($servicePrincipal.DisplayName) (ObjectId: $($servicePrincipal.Id))"
        }

        # Apply the pre-calculated changes
        if ($AssignmentsToAdd.Count -gt 0 -or $AssignmentsToRemove.Count -gt 0) {
            Write-Verbose "Applying pre-calculated changes (Add: $($AssignmentsToAdd.Count), Remove: $($AssignmentsToRemove.Count))"
            Invoke-AppRoleAssignmentChange -ServicePrincipalId $servicePrincipal.Id -AssignmentsToAdd $AssignmentsToAdd -AssignmentsToRemove $AssignmentsToRemove
        }
        else {
            Write-Verbose "No changes to apply"
        }

        return $servicePrincipal
    }
}

function Remove-ServicePrincipal {
    [CmdletBinding(SupportsShouldProcess)]
    param (
        [Parameter(Mandatory = $true)]
        [string]$Id
    )

    if ($PSCmdlet.ShouldProcess("Service Principal with ID $Id", "Remove")) {
        Write-Verbose "Removing service principal with Id: $Id"

        Invoke-RJMgGraphCommand { Remove-MgServicePrincipal -ServicePrincipalId $Id }

        Write-Verbose "Successfully removed service principal: $Id"
    }
}