PSPlumbingRules

0.1.0

Rules/PSAvoidPropertyShapeDrift.psm1

Set-StrictMode -Version Latest

$script:RuleName = 'PSAvoidPropertyShapeDrift'
$script:RuleMessage = 'Avoid changing direct PSCustomObject property shape across switch-driven branches in pipeline-facing commands. Keep emitted object properties stable so pipeline consumers can rely on one shape.'

function Measure-PSAvoidPropertyShapeDrift {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.ScriptBlockAst] $ScriptBlockAst
    )

    $functions = $ScriptBlockAst.FindAll({
        param($Ast)
        $Ast -is [System.Management.Automation.Language.FunctionDefinitionAst]
    }, $true)

    foreach ($functionAst in $functions) {
        if (-not (Test-IsPipelineFacingCommand -FunctionAst $functionAst)) {
            continue
        }

        if (-not (Test-HasPropertyShapeDrift -FunctionAst $functionAst)) {
            continue
        }

        [Microsoft.Windows.PowerShell.ScriptAnalyzer.Generic.DiagnosticRecord]::new(
            $script:RuleMessage,
            $functionAst.Extent,
            $script:RuleName,
            [Microsoft.Windows.PowerShell.ScriptAnalyzer.Generic.DiagnosticSeverity]::Warning,
            $null,
            $null,
            $null
        )
    }
}

function Test-IsPipelineFacingCommand {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.FunctionDefinitionAst] $FunctionAst
    )

    if (-not $FunctionAst.Body.ParamBlock) {
        return $false
    }

    $hasCmdletBinding = $FunctionAst.Body.ParamBlock.Attributes | Where-Object {
        $_.TypeName.GetReflectionType().FullName -eq 'System.Management.Automation.CmdletBindingAttribute'
    }

    if (-not $hasCmdletBinding) {
        return $false
    }

    return $FunctionAst.Body.ProcessBlock -ne $null
}

function Test-HasPropertyShapeDrift {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.FunctionDefinitionAst] $FunctionAst
    )

    $switchParameterNames = @(Get-SwitchParameterNames -FunctionAst $FunctionAst)
    if ($switchParameterNames.Count -eq 0) {
        return $false
    }

    $ifStatements = @(
        $FunctionAst.FindAll({
            param($Ast)
            $Ast -is [System.Management.Automation.Language.IfStatementAst]
        }, $true)
    )

    foreach ($ifStatement in $ifStatements) {
        if (-not (Test-IfStatementUsesSwitchCondition -IfStatementAst $ifStatement -SwitchParameterNames $switchParameterNames)) {
            continue
        }

        $propertySets = @(Get-IfStatementPropertySets -IfStatementAst $ifStatement)
        if ($propertySets.Count -gt 1) {
            return $true
        }
    }

    return $false
}

function Get-SwitchParameterNames {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.FunctionDefinitionAst] $FunctionAst
    )

    @(
        $FunctionAst.Body.ParamBlock.Parameters |
            Where-Object { $_.StaticType -eq [System.Management.Automation.SwitchParameter] } |
            ForEach-Object { $_.Name.VariablePath.UserPath }
    )
}

function Test-IfStatementUsesSwitchCondition {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.IfStatementAst] $IfStatementAst,

        [Parameter(Mandatory)]
        [string[]] $SwitchParameterNames
    )

    foreach ($clause in $IfStatementAst.Clauses) {
        $variables = $clause.Item1.FindAll({
            param($Ast)
            $Ast -is [System.Management.Automation.Language.VariableExpressionAst]
        }, $true)

        foreach ($variable in $variables) {
            if ($SwitchParameterNames -contains $variable.VariablePath.UserPath) {
                return $true
            }
        }
    }

    return $false
}

function Get-IfStatementPropertySets {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.IfStatementAst] $IfStatementAst
    )

    $propertySets = New-Object System.Collections.Generic.List[string]

    foreach ($clause in $IfStatementAst.Clauses) {
        foreach ($propertySet in @(Get-BranchPropertySets -StatementBlockAst $clause.Item2)) {
            if (-not $propertySets.Contains($propertySet)) {
                $propertySets.Add($propertySet)
            }
        }
    }

    if ($IfStatementAst.ElseClause) {
        foreach ($propertySet in @(Get-BranchPropertySets -StatementBlockAst $IfStatementAst.ElseClause)) {
            if (-not $propertySets.Contains($propertySet)) {
                $propertySets.Add($propertySet)
            }
        }
    }
    else {
        foreach ($propertySet in @(Get-ContinuationPropertySets -IfStatementAst $IfStatementAst)) {
            if (-not $propertySets.Contains($propertySet)) {
                $propertySets.Add($propertySet)
            }
        }
    }

    $propertySets
}

function Get-BranchPropertySets {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.Ast] $StatementBlockAst
    )

    $propertySets = New-Object System.Collections.Generic.List[string]

    foreach ($statement in $StatementBlockAst.Statements) {
        $propertySet = Get-PropertySetSignature -StatementAst $statement
        if ($null -eq $propertySet) {
            continue
        }

        if (-not $propertySets.Contains($propertySet)) {
            $propertySets.Add($propertySet)
        }
    }

    $propertySets
}

function Get-ContinuationPropertySets {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.IfStatementAst] $IfStatementAst
    )

    $parentBlock = $IfStatementAst.Parent
    if ($null -eq $parentBlock -or -not ($parentBlock.PSObject.Properties.Name -contains 'Statements')) {
        return @()
    }

    $propertySets = New-Object System.Collections.Generic.List[string]
    $statementIndex = -1

    for ($i = 0; $i -lt $parentBlock.Statements.Count; $i++) {
        if ($parentBlock.Statements[$i] -eq $IfStatementAst) {
            $statementIndex = $i
            break
        }
    }

    if ($statementIndex -lt 0) {
        return @()
    }

    for ($i = $statementIndex + 1; $i -lt $parentBlock.Statements.Count; $i++) {
        $propertySet = Get-PropertySetSignature -StatementAst $parentBlock.Statements[$i]
        if ($null -eq $propertySet) {
            continue
        }

        if (-not $propertySets.Contains($propertySet)) {
            $propertySets.Add($propertySet)
        }
    }

    $propertySets
}

function Get-PropertySetSignature {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.StatementAst] $StatementAst
    )

    $propertySet = @(Get-DirectObjectPropertySet -StatementAst $StatementAst)
    if ($propertySet.Count -eq 0) {
        return $null
    }

    $propertySet -join '|'
}

function Get-DirectObjectPropertySet {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.StatementAst] $StatementAst
    )

    $pipelineAst = $null
    if ($StatementAst -is [System.Management.Automation.Language.PipelineAst]) {
        $pipelineAst = $StatementAst
    }
    elseif ($StatementAst -is [System.Management.Automation.Language.ReturnStatementAst] -and $StatementAst.Pipeline) {
        $pipelineAst = $StatementAst.Pipeline
    }

    if ($null -eq $pipelineAst -or $pipelineAst.PipelineElements.Count -ne 1) {
        return @()
    }

    $element = $pipelineAst.PipelineElements[0]
    if ($element -isnot [System.Management.Automation.Language.CommandExpressionAst]) {
        return @()
    }

    $expression = $element.Expression
    if ($expression -isnot [System.Management.Automation.Language.ConvertExpressionAst]) {
        return @()
    }

    if ($expression.Type.TypeName.FullName -ne 'pscustomobject') {
        return @()
    }

    if ($expression.Child -isnot [System.Management.Automation.Language.HashtableAst]) {
        return @()
    }

    @(
        $expression.Child.KeyValuePairs |
            ForEach-Object { $_.Item1.SafeGetValue() }
    )
}

Export-ModuleMember -Function Measure-PSAvoidPropertyShapeDrift