#!/usr/bin/env bash
set -euo pipefail

# =============================================================================
# setup-exchange-online-certificate.sh
#
# Automates certificate-based authentication setup for Exchange Online:
#   1. Assigns Exchange.ManageAsApp API permission to an app registration
#   2. Assigns the Exchange Administrator role to the app's service principal
#   3. Generates a self-signed certificate (openssl)
#   4. Exports the private key as a passwordless PFX
#   5. Uploads the .cer to the app registration
#   6. Verifies connectivity via pwsh + Connect-ExchangeOnline
#
# Prerequisites: az cli (logged in), openssl, pwsh
# =============================================================================

# ── Configurable variables ───────────────────────────────────────────────────
APP_ID=""
ORGANIZATION=""
CERT_NAME="StorageMonitorExchangeOnline"
CERT_VALIDITY_DAYS=365
OUTPUT_DIR="."
SKIP_VERIFY=false

usage() {
    cat <<EOF
Usage: $0 --app-id <app-client-id> --organization <tenant>.onmicrosoft.com [options]

Required:
  --app-id          Application (client) ID of the Azure app registration
  --organization    Tenant domain (e.g. contoso.onmicrosoft.com)

Options:
  --cert-name       Certificate CN / filename prefix - .cer and .pfx will be appended  (default: StorageMonitorExchangeOnline)
  --validity-days   Certificate validity in days       (default: 365)
  --output-dir      Directory for generated cert files (default: current directory) (must exist)
  --skip-verify     Skip the Exchange Online connectivity check (step 6), removes pwsh dependency
  -h, --help        Show this help message

Prerequisites:
  Tools:            az (Azure CLI), openssl, pwsh (PowerShell Core, not needed with --skip-verify)
  Azure CLI:        Logged in (az login) with permissions to manage App Registrations,
                    grant admin consent, and assign directory roles
  PowerShell:       ExchangeOnlineManagement module (installed automatically if missing)
  Filesystem:       Write permission to the output directory (for .cer and .pfx files)
  App registration: Must already exist with a service principal (Enterprise Application)
EOF
    exit 1
}

# ── Parse arguments ──────────────────────────────────────────────────────────
while [[ $# -gt 0 ]]; do
    case "$1" in
        --app-id)        APP_ID="$2";              shift 2 ;;
        --organization)  ORGANIZATION="$2";        shift 2 ;;
        --cert-name)     CERT_NAME="$2";           shift 2 ;;
        --validity-days) CERT_VALIDITY_DAYS="$2";  shift 2 ;;
        --output-dir)    OUTPUT_DIR="$2";          shift 2 ;;
        --skip-verify)   SKIP_VERIFY=true;         shift ;;
        -h|--help)       usage ;;
        *)               echo "Unknown option: $1"; usage ;;
    esac
done

if [[ -z "$APP_ID" || -z "$ORGANIZATION" ]]; then
    echo "Error: --app-id and --organization are required."
    usage
fi

# ── Preflight checks ────────────────────────────────────────────────────────
for cmd in az openssl; do
    if ! command -v "$cmd" &>/dev/null; then
        echo "Error: '$cmd' is not installed or not on PATH." >&2
        exit 1
    fi
done

if [[ ! -d "$OUTPUT_DIR" ]]; then
    echo "Error: output directory '$OUTPUT_DIR' does not exist." >&2
    exit 1
fi

if [[ "$SKIP_VERIFY" == false ]] && ! command -v pwsh &>/dev/null; then
    echo "Error: 'pwsh' is not installed or not on PATH. Use --skip-verify to skip the connectivity check." >&2
    exit 1
fi

echo "Checking Azure CLI login status..."
az account show --output none 2>/dev/null || {
    echo "Error: not logged in to Azure CLI. Run 'az login' first." >&2
    exit 1
}

# tr -d '\r' strips carriage returns
TENANT_ID=$(az account show --query tenantId -o tsv | tr -d '\r')
echo "Tenant ID: $TENANT_ID"

# ── Resolve the service principal (enterprise app) for this app registration ─
echo ""
echo "=== Resolving service principal for app $APP_ID ==="
# tr -d '\r' strips carriage returns
SP_OBJECT_ID=$(az ad sp show --id "$APP_ID" --query id -o tsv | tr -d '\r') || {
    echo "Error: no service principal found for app '$APP_ID'. Ensure the app registration exists and has a service principal." >&2
    exit 1
}
echo "Service principal object ID: $SP_OBJECT_ID"

# tr -d '\r' strips carriage returns
APP_OBJECT_ID=$(az ad app show --id "$APP_ID" --query id -o tsv | tr -d '\r')
echo "App registration object ID: $APP_OBJECT_ID"

# =============================================================================
# Step 1 — Assign Exchange.ManageAsApp API permission
# =============================================================================
echo ""
echo "=== Step 1: Assigning Exchange.ManageAsApp permission ==="

# Office 365 Exchange Online service principal — well-known app ID
EXCHANGE_RESOURCE_APP_ID="00000002-0000-0ff1-ce00-000000000000"

# Resolve the Exchange Online service principal in this tenant
# tr -d '\r' strips carriage returns
EXCHANGE_SP_ID=$(az ad sp show --id "$EXCHANGE_RESOURCE_APP_ID" --query id -o tsv 2>/dev/null | tr -d '\r' || true)
if [[ -z "$EXCHANGE_SP_ID" ]]; then
    echo "Error: Office 365 Exchange Online service principal not found in tenant." >&2
    exit 1
fi

# Find the Exchange.ManageAsApp app role ID
# tr -d '\r' strips carriage returns
MANAGE_AS_APP_ROLE_ID=$(az ad sp show --id "$EXCHANGE_RESOURCE_APP_ID" \
    --query "appRoles[?value=='Exchange.ManageAsApp'].id | [0]" -o tsv | tr -d '\r')

if [[ -z "$MANAGE_AS_APP_ROLE_ID" ]]; then
    echo "Error: could not find Exchange.ManageAsApp role on Exchange Online SP." >&2
    exit 1
fi
echo "Exchange.ManageAsApp role ID: $MANAGE_AS_APP_ROLE_ID"

# Check if already assigned
# tr -d '\r' strips carriage returns
EXISTING=$(az rest --method GET \
    --uri "https://graph.microsoft.com/v1.0/servicePrincipals/$SP_OBJECT_ID/appRoleAssignments" \
    --query "value[?appRoleId=='$MANAGE_AS_APP_ROLE_ID'] | length(@)" -o tsv 2>/dev/null | tr -d '\r' || echo "0")

if [[ "$EXISTING" -gt 0 ]]; then
    echo "Exchange.ManageAsApp permission is already assigned. Skipping."
else
    echo "Granting Exchange.ManageAsApp..."
    az rest --method POST \
        --uri "https://graph.microsoft.com/v1.0/servicePrincipals/$SP_OBJECT_ID/appRoleAssignments" \
        --headers "Content-Type=application/json" \
        --body "{
            \"principalId\": \"$SP_OBJECT_ID\",
            \"resourceId\": \"$EXCHANGE_SP_ID\",
            \"appRoleId\": \"$MANAGE_AS_APP_ROLE_ID\"
        }" --output none
    echo "Exchange.ManageAsApp permission granted."
fi

# =============================================================================
# Step 2 — Assign Exchange Administrator role
# =============================================================================
echo ""
echo "=== Step 2: Assigning Exchange Administrator directory role ==="

# Get the Exchange Administrator role template ID (well-known)
EXCHANGE_ADMIN_ROLE_TEMPLATE_ID="29232cdf-9323-42fd-ade2-1d097af3e4de"

# Check if already assigned
# tr -d '\r' strips carriage returns
ROLE_ASSIGNED=$(az rest --method GET \
    --uri "https://graph.microsoft.com/v1.0/roleManagement/directory/roleAssignments?\$filter=principalId eq '$SP_OBJECT_ID' and roleDefinitionId eq '$EXCHANGE_ADMIN_ROLE_TEMPLATE_ID'" \
    --query "value | length(@)" -o tsv 2>/dev/null | tr -d '\r' || echo "0")

if [[ "$ROLE_ASSIGNED" -gt 0 ]]; then
    echo "Exchange Administrator role is already assigned. Skipping."
else
    echo "Assigning Exchange Administrator role..."
    az rest --method POST \
        --uri "https://graph.microsoft.com/v1.0/roleManagement/directory/roleAssignments" \
        --headers "Content-Type=application/json" \
        --body "{
            \"@odata.type\": \"#microsoft.graph.unifiedRoleAssignment\",
            \"principalId\": \"$SP_OBJECT_ID\",
            \"roleDefinitionId\": \"$EXCHANGE_ADMIN_ROLE_TEMPLATE_ID\",
            \"directoryScopeId\": \"/\"
        }" --output none
    echo "Exchange Administrator role assigned."
fi

# =============================================================================
# Step 3 — Generate self-signed certificate with openssl
# =============================================================================
echo ""
echo "=== Step 3: Generating self-signed certificate ==="

KEY_FILE="$OUTPUT_DIR/$CERT_NAME.key"
CER_FILE="$OUTPUT_DIR/$CERT_NAME.cer"
PFX_FILE="$OUTPUT_DIR/$CERT_NAME.pfx"

openssl req -x509 -newkey rsa:2048 -sha256 \
    -days "$CERT_VALIDITY_DAYS" \
    -nodes \
    -keyout "$KEY_FILE" \
    -out "$CER_FILE" \
    -subj "/CN=$CERT_NAME" \
    2>/dev/null

echo "Generated: $CER_FILE (public) and $KEY_FILE (private key)"

# =============================================================================
# Step 4 — Export as passwordless PFX
# =============================================================================
echo ""
echo "=== Step 4: Exporting passwordless PFX ==="

openssl pkcs12 -export \
    -out "$PFX_FILE" \
    -inkey "$KEY_FILE" \
    -in "$CER_FILE" \
    -passout pass: \
    2>/dev/null

echo "Generated: $PFX_FILE (no password)"

# Clean up the loose private key file — it's now embedded in the PFX
rm -f "$KEY_FILE"
echo "Removed intermediate key file."

# =============================================================================
# Step 5 — Upload the .cer to the app registration
# =============================================================================
echo ""
echo "=== Step 5: Uploading certificate to app registration ==="

# Upload the .cer to the app registration using az ad app credential reset.
# --append ensures existing credentials are preserved.
az ad app credential reset \
    --id "$APP_ID" \
    --cert "@$CER_FILE" \
    --append \
    --output none

echo "Certificate uploaded."

THUMBPRINT=$(openssl x509 -in "$CER_FILE" -noout -fingerprint -sha1 | sed 's/://g' | cut -d= -f2)
echo "Certificate uploaded. Thumbprint: $THUMBPRINT"

# =============================================================================
# Step 6 — Verify connectivity with pwsh + Connect-ExchangeOnline
# =============================================================================
if [[ "$SKIP_VERIFY" == true ]]; then
    echo ""
    echo "=== Step 6: Skipped (--skip-verify) ==="
else
echo ""
echo "=== Step 6: Verifying Exchange Online connectivity via pwsh ==="

PFX_ABSOLUTE=$(realpath "$PFX_FILE")

# Install the module once, separately from the connection retries
pwsh -NoProfile -NonInteractive -Command "
    if (-not (Get-Module -ListAvailable -Name ExchangeOnlineManagement)) {
        Write-Host 'Installing ExchangeOnlineManagement module...'
        Install-Module -Name ExchangeOnlineManagement -Scope CurrentUser -Force -AllowClobber
    }
    Write-Host 'ExchangeOnlineManagement module ready.'
"

echo "Waiting for certificate to propagate in Azure AD (this can take a few minutes)..."

# Retry up to 4 times with 30s in between (2 mins total)
MAX_RETRIES=4
RETRY_DELAY_SECONDS=30
for (( i=1; i<=MAX_RETRIES; i++ )); do
    if OUTPUT=$(pwsh -NoProfile -NonInteractive -Command "
        \$ErrorActionPreference = 'Stop'
        Import-Module ExchangeOnlineManagement
        Connect-ExchangeOnline -AppId '$APP_ID' -CertificateFilePath '$PFX_ABSOLUTE' -Organization '$ORGANIZATION' -ShowBanner:\$false
        Write-Host 'Connection successful. Running Get-DistributionGroup as a TEST:'
        Get-DistributionGroup | Select-Object -First 5 DisplayName, PrimarySmtpAddress | Format-Table -AutoSize
        Disconnect-ExchangeOnline -Confirm:\$false
        Write-Host 'Disconnected. Verification complete.'
    " 2>&1); then
        echo "$OUTPUT"
        break
    fi

    if [[ $i -eq $MAX_RETRIES ]]; then
        echo "Error: failed to connect after $MAX_RETRIES attempts." >&2
        echo "Removing certificate from app registration..."
        # Look up the credential's key ID by matching the thumbprint (customKeyIdentifier)
        # tr -d '\r' strips carriage returns
        KEY_ID=$(az ad app show --id "$APP_ID" --query "keyCredentials[?customKeyIdentifier=='$THUMBPRINT'].keyId | [0]" -o tsv 2>/dev/null | tr -d '\r')
        if [[ -n "$KEY_ID" ]]; then
            az ad app credential delete --id "$APP_ID" --key-id "$KEY_ID" --cert --output none 2>/dev/null || true
        fi
        rm -f "$PFX_FILE" "$CER_FILE"
        echo "Cleaned up certificate files and app registration credential."
        exit 1
    fi

    echo "Attempt $i/$MAX_RETRIES failed. Retrying in ${RETRY_DELAY_SECONDS}s..."
    sleep "$RETRY_DELAY_SECONDS"
done

fi # end skip-verify check

# =============================================================================
# Summary
# =============================================================================
echo ""
echo "=== Setup complete ==="
echo "  App ID:        $APP_ID"
echo "  Organization:  $ORGANIZATION"
echo "  PFX file:      $(realpath "$PFX_FILE")"
echo "  CER file:      $(realpath "$CER_FILE")"
echo "  Thumbprint:    $THUMBPRINT"
echo ""
echo "Keep the .pfx file secure - it grants access to your Exchange Online tenant without additional credentials."