AzureServiceTags.psm1

$script:PathLocalStore = Join-Path -Path (Split-Path $profile) -ChildPath "lookups"
$script:AzSvcTagCache = @{}
$script:Json = $null
$script:PrefixIndex = @{}

#Region CommonHelperFunctions
function New-LocalStore {
    if (-not (Test-Path -LiteralPath $script:PathLocalStore)) {
        New-Item -Path $script:PathLocalStore -ItemType Directory -Force | Out-Null
    }
}

function Get-JsonFromDiskToMemory {
    param(
        [ValidateSet('Public', 'China', 'AzureGovernment', 'AzureGermany')]
        [string]$Cloud
    )

    New-LocalStore

    $jsonPath = Join-Path $script:PathLocalStore "AzureIPRangesandServiceTags-$Cloud.json"
    if (-not (Test-Path -LiteralPath $jsonPath)) {
        throw "JSON file not found: $jsonPath"
    }

    $script:Json = Get-Content -LiteralPath $jsonPath -Raw | ConvertFrom-Json -Depth 10

    # Write ServiceTags
    $script:Json.values | Where-Object { $_.Name -notmatch '\.' } | Select-Object -ExpandProperty Name | Sort-Object -Unique | Set-Content (Join-Path $script:PathLocalStore "AzureServiceTags-$Cloud.txt") -Force

    # Write Regions
    $script:Json.values.properties.region | Where-Object { $_ } | Sort-Object -Unique | Set-Content (Join-Path $script:PathLocalStore "AzureRegions-$Cloud.txt") -Force

    # JSON changed -> invalidate prefix index for this cloud
    if ($script:PrefixIndex.ContainsKey($Cloud)) {
        $script:PrefixIndex.Remove($Cloud) | Out-Null
    }
}

function Get-JsonFromInternetToDisk {
    param(
        [ValidateSet('Public', 'China', 'AzureGovernment', 'AzureGermany')]
        [string]$Cloud
    )

    New-LocalStore

    $uri = switch ($Cloud) {
        'China' { 'https://www.microsoft.com/en-us/download/details.aspx?id=57062' }
        'AzureGovernment' { 'https://www.microsoft.com/en-us/download/details.aspx?id=57063' }
        'AzureGermany' { 'https://www.microsoft.com/en-us/download/details.aspx?id=57064' }
        default { 'https://www.microsoft.com/en-us/download/details.aspx?id=56519' }
    }

    $filePath = Join-Path $script:PathLocalStore "AzureIPRangesandServiceTags-$Cloud.json"

    $download = $false
    if (-not (Test-Path -LiteralPath $filePath)) {
        $download = $true
    }
    else {
        $age = (Get-Date) - (Get-Item -LiteralPath $filePath).LastWriteTime
        if ($age.TotalDays -gt 2) { $download = $true }
    }

    if ($download) {
        $resp = Invoke-WebRequest -Uri $uri -UseBasicParsing

        # Prefer matching the JSON link directly (avoids PropertyNotFoundException on .download)
        $downloadUrl = $resp.Links |
            Where-Object { $_.href -match 'AzureIPRangesandServiceTags' -and $_.href -match '\.json(\?|$)' } |
            Select-Object -ExpandProperty href -First 1

        # Fallback: if page structure changes, check 'download' property safely
        if (-not $downloadUrl) {
            $downloadUrl = $resp.Links | Where-Object { $_.PSObject.Properties.Match('download').Count -gt 0 } | Select-Object -ExpandProperty href -First 1
        }

        if (-not $downloadUrl) {
            throw "Could not find the JSON download link on: $uri"
        }

        Invoke-WebRequest -Uri $downloadUrl -OutFile $filePath -UseBasicParsing
    }

    # Load JSON if not loaded or cloud mismatch
    if (-not $script:Json -or ($script:Json.cloud -and $script:Json.cloud -ne $Cloud)) {
        Get-JsonFromDiskToMemory -Cloud $Cloud
    }
}

function Get-JsonData {
    param(
        [ValidateSet('Public', 'China', 'AzureGovernment', 'AzureGermany')]
        [string]$Cloud
    )

    Get-JsonFromInternetToDisk -Cloud $Cloud | Out-Null

    if (-not $script:AzSvcTagCache.ContainsKey($Cloud)) {
        $script:AzSvcTagCache[$Cloud] = @{
            ServiceTags = Get-Content (Join-Path $script:PathLocalStore "AzureServiceTags-$Cloud.txt")
            Regions     = Get-Content (Join-Path $script:PathLocalStore "AzureRegions-$Cloud.txt")
        }
    }
}

function Get-CloudFromCompleter {
    param($fakeBoundParameters)
    if ($fakeBoundParameters.ContainsKey('Cloud') -and $fakeBoundParameters['Cloud']) {
        return [string]$fakeBoundParameters['Cloud']
    }
    'Public'
}

function Complete-List {
    param($List, $Word)
    $List | Where-Object { $_ -like "$Word*" }
}
#EndRegion CommonHelperFunctions

#Region AddressPrefixHelperFunctions
function Convert-IPToBigInt {
    param(
        [Parameter(Mandatory)]
        [System.Net.IPAddress]$IPAddress
    )
    # BigInteger expects little-endian; IP bytes are big-endian
    $bytes = $IPAddress.GetAddressBytes()
    [Array]::Reverse($bytes)

    # Add a 0 byte to force unsigned positive BigInteger
    $unsigned = New-Object byte[] ($bytes.Length + 1)
    [Array]::Copy($bytes, 0, $unsigned, 0, $bytes.Length)

    return [System.Numerics.BigInteger]::new($unsigned)
}

function ConvertFrom-IpOrCidrToRange {
    <#
      Returns:
        [pscustomobject]@{
          Family = 4|6
          Bits = 32|128
          Start = BigInteger
          End = BigInteger
          Input = original string
          PrefixLength = int? (null for single IP)
        }
    #>

    param(
        [Parameter(Mandatory)]
        [string]$Text
    )

    $t = $Text.Trim()

    $ipPart = $t
    $prefixLen = $null

    if ($t -match '^(?<ip>.+?)/(?<plen>\d+)$') {
        $ipPart = $Matches.ip
        $prefixLen = [int]$Matches.plen
    }

    $ip = $null
    if (-not [System.Net.IPAddress]::TryParse($ipPart, [ref]$ip)) {
        throw "Invalid IP/CIDR: '$Text'"
    }

    $isV4 = ($ip.AddressFamily -eq [System.Net.Sockets.AddressFamily]::InterNetwork)
    $isV6 = ($ip.AddressFamily -eq [System.Net.Sockets.AddressFamily]::InterNetworkV6)
    if (-not ($isV4 -or $isV6)) {
        throw "Unsupported IP family for '$Text'"
    }

    $bits = if ($isV4) { 32 } else { 128 }

    if ($null -ne $prefixLen) {
        if ($prefixLen -lt 0 -or $prefixLen -gt $bits) {
            throw "Invalid prefix length '$prefixLen' for '$Text'"
        }
    }

    $ipInt = Convert-IPToBigInt -IPAddress $ip

    if ($null -eq $prefixLen) {
        return [pscustomobject]@{
            Family       = if ($isV4) { 4 } else { 6 }
            Bits         = $bits
            Start        = $ipInt
            End          = $ipInt
            Input        = $Text
            PrefixLength = $null
        }
    }

    # Build mask and compute network range
    $hostBits = $bits - $prefixLen

    # mask = ((2^bits - 1) << hostBits) & (2^bits - 1)
    $allOnes = ([System.Numerics.BigInteger]::One -shl $bits) - 1
    $mask = (($allOnes -shl $hostBits) -band $allOnes)

    $start = $ipInt -band $mask
    $end = $start + (([System.Numerics.BigInteger]::One -shl $hostBits) - 1)

    return [pscustomobject]@{
        Family       = if ($isV4) { 4 } else { 6 }
        Bits         = $bits
        Start        = $start
        End          = $end
        Input        = $Text
        PrefixLength = $prefixLen
    }
}

function New-PrefixIndex {
    param(
        [ValidateSet('Public', 'China', 'AzureGovernment', 'AzureGermany')]
        [string]$Cloud
    )

    if ($script:PrefixIndex.ContainsKey($Cloud)) { return }

    # Ensure JSON loaded
    Get-JsonFromInternetToDisk -Cloud $Cloud | Out-Null

    $v4 = New-Object System.Collections.Generic.List[object]
    $v6 = New-Object System.Collections.Generic.List[object]

    foreach ($val in $script:Json.values) {
        $p = $val.properties
        if (-not $p) { continue }
        if (-not $p.addressPrefixes) { continue }

        foreach ($prefix in $p.addressPrefixes) {
            if (-not $prefix) { continue }

            try {
                $r = ConvertFrom-IpOrCidrToRange -Text ([string]$prefix)

                $row = [pscustomobject]@{
                    Cloud         = $Cloud
                    Platform      = $p.platform
                    SystemService = $p.systemService
                    Region        = $p.region
                    RegionId      = $p.regionId
                    AddressPrefix = [string]$prefix
                    Start         = $r.Start
                    End           = $r.End
                    Family        = $r.Family
                    Bits          = $r.Bits
                    PrefixLength  = $r.PrefixLength
                }

                if ($r.Family -eq 4) { $v4.Add($row) } else { $v6.Add($row) }
            }
            catch {
                # If a prefix is malformed (rare), ignore it
                continue
            }
        }
    }

    $script:PrefixIndex[$Cloud] = @{
        v4 = $v4
        v6 = $v6
    }
}
#EndRegion AddressPrefixHelperFunctions

#Region Functions
function Get-AzureServiceTag {
    [CmdletBinding()]
    param(
        [ValidateSet('Public', 'China', 'AzureGovernment', 'AzureGermany')]
        [string]$Cloud = 'Public',

        [string]$Name,
        [string]$Region
    )

    Get-JsonFromInternetToDisk -Cloud $Cloud

    $script:Json.Values.properties | Where-Object { ($Name -and $_.systemService -match $Name -or -not $Name) -and ($Region -and $_.region -match $Region -or -not $Region) }
}

function Get-AzureIpRange {
    <#
    .SYNOPSIS
      Look up Azure region/systemService by IP or CIDR and find matching addressPrefixes.
 
    .PARAMETER Address
      One or more IPs/CIDRs (IPv4/IPv6). Examples:
        - 51.120.224.216
        - 51.120.224.216/29
        - 2603:1020:f04::340
        - 2603:1020:f04::340/123
 
    .PARAMETER Cloud
      Public/China/AzureGovernment/AzureGermany
 
    .PARAMETER Match
      - ContainedOnly: input range must be fully inside the prefix range
      - Overlap: any intersection counts (default, usually what people want)
    #>

    [CmdletBinding()]
    param(
        [Parameter(Mandatory, ValueFromPipeline, ValueFromPipelineByPropertyName)]
        [Alias('IP', 'IPs', 'Prefix', 'CIDR')]
        [string[]]$Address,

        [ValidateSet('Public', 'China', 'AzureGovernment', 'AzureGermany')]
        [string]$Cloud = 'Public',

        [ValidateSet('Overlap', 'ContainedOnly')]
        [string]$Match = 'Overlap'
    )

    begin {
        New-PrefixIndex -Cloud $Cloud
        $idx = $script:PrefixIndex[$Cloud]
    }

    process {
        foreach ($a in $Address) {
            $query = ConvertFrom-IpOrCidrToRange -Text $a

            $list = if ($query.Family -eq 4) { $idx.v4 } else { $idx.v6 }

            foreach ($p in $list) {
                # Compare BigInteger ranges
                $contained = ($query.Start -ge $p.Start) -and ($query.End -le $p.End)
                $overlap = ($query.Start -le $p.End) -and ($query.End -ge $p.Start)

                if ($Match -eq 'ContainedOnly') {
                    if (-not $contained) { continue }
                    $matchType = 'Contained'
                }
                else {
                    if (-not $overlap) { continue }
                    $matchType = if ($contained) { 'Contained' } else { 'Overlap' }
                }

                [pscustomobject]@{
                    Input         = $query.Input
                    InputFamily   = if ($query.Family -eq 4) { 'IPv4' } else { 'IPv6' }
                    InputPrefix   = $query.PrefixLength
                    Cloud         = $p.Cloud
                    Platform      = $p.Platform
                    SystemService = $p.SystemService
                    Region        = $p.Region
                    RegionId      = $p.RegionId
                    AddressPrefix = $p.AddressPrefix
                    MatchType     = $matchType
                }
            }
        }
    }
}
#EndRegion Functions

#Region ArgumentCompleters
Register-ArgumentCompleter -CommandName Get-AzureServiceTag -ParameterName Name -ScriptBlock {
    param($cmd, $param, $word, $ast, $fake)
    $cloud = Get-CloudFromCompleter $fake
    Get-JsonData -Cloud $cloud
    Complete-List $script:AzSvcTagCache[$cloud].ServiceTags $word
}

Register-ArgumentCompleter -CommandName Get-AzureServiceTag -ParameterName Region -ScriptBlock {
    param($cmd, $param, $word, $ast, $fake)
    $cloud = Get-CloudFromCompleter $fake
    Get-JsonData -Cloud $cloud
    Complete-List $script:AzSvcTagCache[$cloud].Regions $word
}
#EndRegion ArgumentCompleters

Export-ModuleMember -Function Get-AzureServiceTag, Get-AzureIpRange