<#
.SYNOPSIS
  CyberHoot M365 configuration helper:
  - Advanced Delivery (PhishSimOverridePolicy + ExoPhishSimOverrideRule) merge-safe apply/rollback
  - Safe Links DoNotRewriteUrls merge-safe apply/rollback
  - Pre-flight checks + JSON snapshots
  - Validate mode + optional verbose logging + summary

.PARAMETER Mode
  Apply, Rollback, or Validate

.PARAMETER OutputDir
  Directory for JSON snapshots (default: current directory)

.PARAMETER Verbose
  Enables verbose logging (Write-Verbose output)

.NOTES
  Run in PowerShell (pwsh) on Windows/macOS/Linux.
  Requires ExchangeOnlineManagement module and sufficient permissions.
#>

param(
  [ValidateSet("Apply","Rollback","Validate")]
  [string]$Mode = "Apply",

  [string]$OutputDir = ".",

  [switch]$Verbose
)

if ($Verbose) { $VerbosePreference = 'Continue' }

# ---------------------------
# CONFIG (CyberHoot)
# ---------------------------
$SendingDomains = @(
  "cyberhoot.com",
  "docunotice.com",
  "messagecenters.net",
  "securedinbox.net",
  "notificationhub.net"
)

$SenderIpRanges = @(
  "23.20.251.170/32", "52.7.191.238/32", "52.6.6.155/32", "18.213.175.22/32"
)

# Advanced Delivery URL patterns for phishing sim overrides
$SimulationUrls = @(
  "*.cyberhoot.com/*",
  "*.docunotice.com/*",
  "*.messagecenters.net/*",
  "*.securedinbox.net/*",
  "*.notificationhub.net/*"
)

# Safe Links exclusions: use plain domains/URLs (avoid wildcard syntax here)
$DoNotRewriteUrls = @(
  "cyberhoot.com",
  "docunotice.com",
  "messagecenters.net",
  "securedinbox.net",
  "notificationhub.net"
)

# Names we create/manage
$PhishSimPolicyName   = "PhishSimOverridePolicy"
$PhishSimRuleName     = "CyberHoot - PhishSim Override Rule"
$SafeLinksPolicyName  = "CyberHoot - Safe Links Exclusions"
$SafeLinksRuleName    = "CyberHoot - Safe Links Exclusions Rule"

# ---------------------------
# Helpers
# ---------------------------
function Ensure-OutputDir {
  param([string]$Dir)
  if (-not (Test-Path -Path $Dir)) {
    New-Item -ItemType Directory -Path $Dir | Out-Null
  }
}

function Write-Snapshot {
  param(
    [string]$Phase,
    [object]$Data
  )
  Ensure-OutputDir -Dir $OutputDir
  $ts = Get-Date -Format "yyyyMMdd-HHmmss"
  $path = Join-Path $OutputDir "cyberhoot-m365-$Mode-$Phase-$ts.json"
  $Data | ConvertTo-Json -Depth 10 | Out-File -FilePath $path -Encoding utf8
  Write-Verbose "Snapshot saved: $path"
  return $path
}

function Get-ArraySafe {
  param([object]$Value)
  if ($null -eq $Value) { return @() }
  return @($Value)
}

function Merge-Unique {
  param([object[]]$A, [object[]]$B)
  return @($A + $B | Where-Object { $_ -ne $null -and $_ -ne "" } | Sort-Object -Unique)
}

function Remove-Items {
  param([object[]]$Source, [object[]]$ToRemove)
  $removeSet = New-Object 'System.Collections.Generic.HashSet[string]'
  foreach ($r in $ToRemove) { [void]$removeSet.Add([string]$r) }

  $out = @()
  foreach ($s in $Source) {
    if (-not $removeSet.Contains([string]$s)) { $out += $s }
  }
  return $out
}

function Require-Cmdlets {
  param([string[]]$Names)
  $missing = @()
  foreach ($n in $Names) {
    if (-not (Get-Command $n -ErrorAction SilentlyContinue)) { $missing += $n }
  }
  if ($missing.Count -gt 0) {
    Write-Error ("Missing required cmdlets: " + ($missing -join ", ") + ". Ensure ExchangeOnlineManagement is installed/updated and you have required permissions/licensing.")
    exit 1
  }
}

function Get-State {
  $state = [ordered]@{}
  $state.AcceptedDomains = (Get-AcceptedDomain).Name

  $state.SafeLinksPolicy = Get-SafeLinksPolicy -Identity $SafeLinksPolicyName -ErrorAction SilentlyContinue
  $state.SafeLinksRule   = Get-SafeLinksRule   -Identity $SafeLinksRuleName   -ErrorAction SilentlyContinue

  $state.PhishSimPolicy  = Get-PhishSimOverridePolicy -ErrorAction SilentlyContinue
  $state.PhishSimRule    = Get-ExoPhishSimOverrideRule -Identity $PhishSimRuleName -ErrorAction SilentlyContinue

  return $state
}

function Compare-Desired {
  param([object]$State)

  $report = [ordered]@{}

  $psRule = $State.PhishSimRule
  $slPol  = $State.SafeLinksPolicy
  $slRule = $State.SafeLinksRule

  $psDomains = @(); $psIPs=@(); $psUrls=@()
  if ($psRule) {
    $psDomains = Get-ArraySafe $psRule.Domains
    $psIPs     = Get-ArraySafe $psRule.SenderIpRanges
    $psUrls    = Get-ArraySafe $psRule.PhishSimOverrideUrl
  }

  $slDnr = @()
  if ($slPol) { $slDnr = Get-ArraySafe $slPol.DoNotRewriteUrls }

  $report.AdvancedDelivery = [ordered]@{
    RuleExists            = [bool]$psRule
    MissingDomains        = @($SendingDomains | Where-Object { $_ -notin $psDomains })
    MissingSenderIpRanges = @($SenderIpRanges | Where-Object { $_ -notin $psIPs })
    MissingSimulationUrls = @($SimulationUrls | Where-Object { $_ -notin $psUrls })
  }

  $report.SafeLinks = [ordered]@{
    PolicyExists        = [bool]$slPol
    RuleExists          = [bool]$slRule
    MissingDoNotRewrite = @($DoNotRewriteUrls | Where-Object { $_ -notin $slDnr })
  }

  return $report
}

# ---------------------------
# Pre-flight
# ---------------------------
Write-Host "Mode: $Mode"
Write-Verbose "PowerShell version: $($PSVersionTable.PSVersion)"

if (-not (Get-Module -ListAvailable -Name ExchangeOnlineManagement)) {
  Write-Host "Installing ExchangeOnlineManagement..."
  Install-Module -Name ExchangeOnlineManagement -Force
}
Import-Module ExchangeOnlineManagement

Require-Cmdlets -Names @("Connect-ExchangeOnline","Disconnect-ExchangeOnline")

$AdminUpn = Read-Host "Enter your Exchange admin email or UPN"
Connect-ExchangeOnline -UserPrincipalName $AdminUpn

Require-Cmdlets -Names @(
  "Get-AcceptedDomain",
  "Get-SafeLinksPolicy","Set-SafeLinksPolicy","New-SafeLinksPolicy",
  "Get-SafeLinksRule","Set-SafeLinksRule","New-SafeLinksRule",
  "Get-PhishSimOverridePolicy","Set-PhishSimOverridePolicy","New-PhishSimOverridePolicy",
  "Get-ExoPhishSimOverrideRule","Set-ExoPhishSimOverrideRule","New-ExoPhishSimOverrideRule"
)

try {
  [void](Get-AcceptedDomain -ErrorAction Stop)
  [void](Get-PhishSimOverridePolicy -ErrorAction Stop)
  [void](Get-SafeLinksPolicy -ErrorAction Stop)
} catch {
  Write-Error "Pre-flight read check failed. Details: $_"
  Disconnect-ExchangeOnline -Confirm:$false
  exit 1
}

Write-Host "Pre-flight checks passed."

# ---------------------------
# Snapshot BEFORE
# ---------------------------
$before = $null
try { $before = Get-State } catch { $before = [ordered]@{ Error="$_" } }
$beforeSnap = Write-Snapshot -Phase "before" -Data $before

# ---------------------------
# VALIDATE (no changes)
# ---------------------------
if ($Mode -eq "Validate") {
  $state  = Get-State
  $drift  = Compare-Desired -State $state

  Write-Host "`nValidation report:"
  Write-Host "Advanced Delivery:"
  Write-Host "  Rule exists: $($drift.AdvancedDelivery.RuleExists)"
  Write-Host "  Missing domains: $($drift.AdvancedDelivery.MissingDomains.Count)"
  Write-Host "  Missing IP ranges: $($drift.AdvancedDelivery.MissingSenderIpRanges.Count)"
  Write-Host "  Missing simulation URLs: $($drift.AdvancedDelivery.MissingSimulationUrls.Count)"

  Write-Host "Safe Links:"
  Write-Host "  Policy exists: $($drift.SafeLinks.PolicyExists)"
  Write-Host "  Rule exists: $($drift.SafeLinks.RuleExists)"
  Write-Host "  Missing DoNotRewriteUrls: $($drift.SafeLinks.MissingDoNotRewrite.Count)"

  $afterSnap = Write-Snapshot -Phase "after" -Data $state

  Write-Host "`nSummary:"
  Write-Host "  Mode: Validate"
  Write-Host "  Before snapshot: $beforeSnap"
  Write-Host "  After snapshot:  $afterSnap"
  Write-Host "  Domains desired: $($SendingDomains.Count)"
  Write-Host "  IPs desired: $($SenderIpRanges.Count)"
  Write-Host "  Simulation URLs desired: $($SimulationUrls.Count)"
  Write-Host "  DoNotRewriteUrls desired: $($DoNotRewriteUrls.Count)"

  Disconnect-ExchangeOnline -Confirm:$false
  Write-Host "Disconnected."
  exit 0
}

# ---------------------------
# APPLY
# ---------------------------
$summary = [ordered]@{
  Mode = $Mode
  DomainsAdded = 0
  IpsAdded = 0
  SimulationUrlsAdded = 0
  DoNotRewriteAdded = 0
  DomainsRemoved = 0
  IpsRemoved = 0
  SimulationUrlsRemoved = 0
  DoNotRewriteRemoved = 0
}

if ($Mode -eq "Apply") {

  # 1) Advanced Delivery: ensure policy enabled
  $policy = Get-PhishSimOverridePolicy -ErrorAction SilentlyContinue
  if (-not $policy) {
    Write-Verbose "Creating PhishSimOverridePolicy: $PhishSimPolicyName"
    New-PhishSimOverridePolicy -Name $PhishSimPolicyName | Out-Null
    Set-PhishSimOverridePolicy -Identity $PhishSimPolicyName -Enabled $true | Out-Null
    $policyIdentity = $PhishSimPolicyName
  } else {
    Write-Verbose "Enabling existing PhishSimOverridePolicy: $($policy.Identity)"
    Set-PhishSimOverridePolicy -Identity $policy.Identity -Enabled $true | Out-Null
    $policyIdentity = $policy.Identity
  }

  # 1b) Rule create/merge-update
  $existingRule = Get-ExoPhishSimOverrideRule -Identity $PhishSimRuleName -ErrorAction SilentlyContinue
  if (-not $existingRule) {
    Write-Verbose "Creating ExoPhishSimOverrideRule: $PhishSimRuleName"
    New-ExoPhishSimOverrideRule -Name $PhishSimRuleName `
      -Policy $policyIdentity `
      -Domains $SendingDomains `
      -SenderIpRanges $SenderIpRanges `
      -PhishSimOverrideUrl $SimulationUrls | Out-Null

    $summary.DomainsAdded = $SendingDomains.Count
    $summary.IpsAdded = $SenderIpRanges.Count
    $summary.SimulationUrlsAdded = $SimulationUrls.Count
  } else {
    $existingDomains = Get-ArraySafe $existingRule.Domains
    $existingIPs     = Get-ArraySafe $existingRule.SenderIpRanges
    $existingUrls    = Get-ArraySafe $existingRule.PhishSimOverrideUrl

    $mergedDomains = Merge-Unique $existingDomains $SendingDomains
    $mergedIPs     = Merge-Unique $existingIPs     $SenderIpRanges
    $mergedUrls    = Merge-Unique $existingUrls    $SimulationUrls

    $summary.DomainsAdded = @($SendingDomains | Where-Object { $_ -notin $existingDomains }).Count
    $summary.IpsAdded     = @($SenderIpRanges | Where-Object { $_ -notin $existingIPs }).Count
    $summary.SimulationUrlsAdded = @($SimulationUrls | Where-Object { $_ -notin $existingUrls }).Count

    Write-Verbose "Updating ExoPhishSimOverrideRule (merge-safe): $PhishSimRuleName"
    Set-ExoPhishSimOverrideRule -Identity $PhishSimRuleName `
      -Domains $mergedDomains `
      -SenderIpRanges $mergedIPs `
      -PhishSimOverrideUrl $mergedUrls | Out-Null
  }

  # 2) Safe Links: create/merge-update policy
  $existingSlPolicy = Get-SafeLinksPolicy -Identity $SafeLinksPolicyName -ErrorAction SilentlyContinue
  if (-not $existingSlPolicy) {
    Write-Verbose "Creating Safe Links policy: $SafeLinksPolicyName"
    New-SafeLinksPolicy -Name $SafeLinksPolicyName `
      -IsEnabled $true `
      -EnableSafeLinksForEmail $true `
      -ScanUrls $true `
      -TrackClicks $true `
      -EnableForInternalSenders $true `
      -DoNotRewriteUrls $DoNotRewriteUrls | Out-Null

    $summary.DoNotRewriteAdded = $DoNotRewriteUrls.Count
  } else {
    $current = Get-ArraySafe $existingSlPolicy.DoNotRewriteUrls
    $merged  = Merge-Unique $current $DoNotRewriteUrls

    $summary.DoNotRewriteAdded = @($DoNotRewriteUrls | Where-Object { $_ -notin $current }).Count

    Write-Verbose "Updating Safe Links policy (merge-safe): $SafeLinksPolicyName"
    Set-SafeLinksPolicy -Identity $SafeLinksPolicyName `
      -DoNotRewriteUrls $merged `
      -IsEnabled $true `
      -EnableSafeLinksForEmail $true `
      -ScanUrls $true `
      -TrackClicks $true `
      -EnableForInternalSenders $true | Out-Null
  }

  # 2b) Safe Links rule tenant-wide
  $acceptedDomains = (Get-AcceptedDomain).Name
  $existingSlRule = Get-SafeLinksRule -Identity $SafeLinksRuleName -ErrorAction SilentlyContinue
  if (-not $existingSlRule) {
    Write-Verbose "Creating Safe Links rule: $SafeLinksRuleName"
    New-SafeLinksRule -Name $SafeLinksRuleName `
      -SafeLinksPolicy $SafeLinksPolicyName `
      -RecipientDomainIs $acceptedDomains `
      -Priority 0 `
      -Enabled $true | Out-Null
  } else {
    Write-Verbose "Updating Safe Links rule: $SafeLinksRuleName"
    Set-SafeLinksRule -Identity $SafeLinksRuleName `
      -SafeLinksPolicy $SafeLinksPolicyName `
      -RecipientDomainIs $acceptedDomains | Out-Null
  }

  Write-Host "Apply completed."
}

# ---------------------------
# ROLLBACK
# ---------------------------
if ($Mode -eq "Rollback") {

  # 1) Advanced Delivery: remove only CyberHoot items from our CyberHoot rule
  $existingRule = Get-ExoPhishSimOverrideRule -Identity $PhishSimRuleName -ErrorAction SilentlyContinue
  if ($existingRule) {
    $existingDomains = Get-ArraySafe $existingRule.Domains
    $existingIPs     = Get-ArraySafe $existingRule.SenderIpRanges
    $existingUrls    = Get-ArraySafe $existingRule.PhishSimOverrideUrl

    $newDomains = Remove-Items -Source $existingDomains -ToRemove $SendingDomains
    $newIPs     = Remove-Items -Source $existingIPs     -ToRemove $SenderIpRanges
    $newUrls    = Remove-Items -Source $existingUrls    -ToRemove $SimulationUrls

    $summary.DomainsRemoved = @($SendingDomains | Where-Object { $_ -in $existingDomains }).Count
    $summary.IpsRemoved     = @($SenderIpRanges | Where-Object { $_ -in $existingIPs }).Count
    $summary.SimulationUrlsRemoved = @($SimulationUrls | Where-Object { $_ -in $existingUrls }).Count

    if (($newDomains.Count + $newIPs.Count + $newUrls.Count) -eq 0) {
      if (Get-Command Remove-ExoPhishSimOverrideRule -ErrorAction SilentlyContinue) {
        Remove-ExoPhishSimOverrideRule -Identity $PhishSimRuleName -Confirm:$false
      } else {
        Set-ExoPhishSimOverrideRule -Identity $PhishSimRuleName -Domains @() -SenderIpRanges @() -PhishSimOverrideUrl @() | Out-Null
      }
      Write-Host "Removed PhishSim override rule: $PhishSimRuleName"
    } else {
      Set-ExoPhishSimOverrideRule -Identity $PhishSimRuleName `
        -Domains $newDomains `
        -SenderIpRanges $newIPs `
        -PhishSimOverrideUrl $newUrls | Out-Null
      Write-Host "Updated PhishSim override rule (CyberHoot entries removed): $PhishSimRuleName"
    }
  } else {
    Write-Host "PhishSim override rule not found: $PhishSimRuleName"
  }

  # 2) Safe Links: remove only CyberHoot DoNotRewriteUrls from our CyberHoot policy
  $existingSlPolicy = Get-SafeLinksPolicy -Identity $SafeLinksPolicyName -ErrorAction SilentlyContinue
  if ($existingSlPolicy) {
    $current = Get-ArraySafe $existingSlPolicy.DoNotRewriteUrls
    $newList = Remove-Items -Source $current -ToRemove $DoNotRewriteUrls

    $summary.DoNotRewriteRemoved = @($DoNotRewriteUrls | Where-Object { $_ -in $current }).Count

    Set-SafeLinksPolicy -Identity $SafeLinksPolicyName -DoNotRewriteUrls $newList | Out-Null
    Write-Host "Updated Safe Links policy (CyberHoot exclusions removed): $SafeLinksPolicyName"

    # Remove rule/policy if empty (best-effort)
    $existingSlRule = Get-SafeLinksRule -Identity $SafeLinksRuleName -ErrorAction SilentlyContinue
    if ($existingSlRule -and (Get-Command Remove-SafeLinksRule -ErrorAction SilentlyContinue)) {
      Remove-SafeLinksRule -Identity $SafeLinksRuleName -Confirm:$false
      Write-Host "Removed Safe Links rule: $SafeLinksRuleName"
    }

    if ($newList.Count -eq 0 -and (Get-Command Remove-SafeLinksPolicy -ErrorAction SilentlyContinue)) {
      Remove-SafeLinksPolicy -Identity $SafeLinksPolicyName -Confirm:$false
      Write-Host "Removed Safe Links policy (now empty): $SafeLinksPolicyName"
    }
  } else {
    Write-Host "Safe Links policy not found: $SafeLinksPolicyName"
  }

  Write-Host "Rollback completed."
}

# ---------------------------
# Snapshot AFTER + Summary
# ---------------------------
$after = $null
try { $after = Get-State } catch { $after = [ordered]@{ Error="$_" } }
$afterSnap = Write-Snapshot -Phase "after" -Data $after

Write-Host "`nSummary:"
Write-Host "  Mode: $Mode"
Write-Host "  Before snapshot: $beforeSnap"
Write-Host "  After snapshot:  $afterSnap"

if ($Mode -eq "Apply") {
  Write-Host "  Domains added:            $($summary.DomainsAdded)"
  Write-Host "  IPs added:                $($summary.IpsAdded)"
  Write-Host "  Simulation URLs added:    $($summary.SimulationUrlsAdded)"
  Write-Host "  DoNotRewriteUrls added:   $($summary.DoNotRewriteAdded)"
}

if ($Mode -eq "Rollback") {
  Write-Host "  Domains removed:          $($summary.DomainsRemoved)"
  Write-Host "  IPs removed:              $($summary.IpsRemoved)"
  Write-Host "  Simulation URLs removed:  $($summary.SimulationUrlsRemoved)"
  Write-Host "  DoNotRewriteUrls removed: $($summary.DoNotRewriteRemoved)"
}

Disconnect-ExchangeOnline -Confirm:$false
Write-Host "Disconnected."