PowerShell Scripts with AST: Advanced Analysis Techniques

In the previous part, we attempted to create a tool that lists all .NET methods used in a script. While this worked well using the tokenizer, it quickly became clear that certain script information—like a list of .NET methods—requires additional contextual details to be truly useful. For example, it would be beneficial to know more about the origin of the .NET methods.

In this part, we will use the AST (Abstract Syntax Tree) to gain deeper insights into the script code.

Asking AST for Information

In the previous part, we already explored how to access the AST for a particular script:

# path to a PowerShell script to analyze (make sure it exists!)
$path = "C:\test\some_file.ps1"
# empty variables, these must exist. The method returns syntax errors and
# tokens in these variables (by reference) after the call
$syntaxErrors = $null 
$tokens = $null
# parse the file
$ast = [System.Management.Automation.Language.Parser]::ParseFile($Path, [ref] $tokens, [ref]$syntaxErrors)

You’ve already seen that you need to actively query the AST for information, and we used this line to display all the information the AST knows about your script:

$ast.FindAll({$true}, $true) | ForEach-Object { $_.GetType().Name } 

The result is a list of type names, and now it's your task to identify the type that describes what you're looking for. When you do this for the first time, it might feel awkward.

However, the AST types are always the same, so once you know that an InvokeMemberExpressionAst represents invoking a .NET method, you can use this type whenever you want to identify a .NET method call. If you're not interested in .NET methods, there's almost certainly another AST type that describes something else you might find more exciting, such as a VariableExpressionAst or a CommandExpressionAst.

Asking AST for Specific Information

Once you've identified an AST type that interests you, it's time to ask the AST a specific question. Previously, we asked the AST to display everything:

$ast.FindAll({$true}, $true)

• { $true } is a filter statement that is executed for each AST object found in the script you're examining. When it returns $true, the corresponding AST object is returned.
• $true instructs the AST to analyze nested script blocks as well (similar to how recursion works with Get-ChildItem). If you specify $false instead, all nested code blocks are ignored.

To see all .NET member invocations throughout the script, use the following specific filter statement:

$path = " C:\test\some_file.ps1"

# get AST for script file
$errors = 
$tokens = $null
$ast = [System.Management.Automation.Language.Parser]::ParseFile($Path, [ref] $tokens, [ref]$errors)

# what would you like to know about this script?
$filter = { 
  param($astObject)  
  
  $astObject -is [System.Management.Automation.Language.InvokeMemberExpressionAst] 
}

# search for all AST objects, including nested script blocks
$recurse = $true
   
# find the AST object
$ast.FindAll($filter, $recurse) 

Provided that the script specified in $path actually uses .NET method calls, the result will look similar to this:

Arguments  : 
Expression : [System.Net.Sockets.TcpClient]
Member     : new
Static     : True
StaticType : System.Object
Extent     : [System.Net.Sockets.TcpClient]::new()
Parent     : [System.Net.Sockets.TcpClient]::new()

Arguments  : {$ComputerName, $Port}
Expression : $client
Member     : ConnectAsync
Static     : False
StaticType : System.Object
Extent     : $client.ConnectAsync($ComputerName, $Port)
Parent     : $client.ConnectAsync($ComputerName, $Port)

Arguments  : {$TimeoutMilliSec}
Expression : $task
Member     : Wait
Static     : False
StaticType : System.Object
Extent     : $task.Wait($TimeoutMilliSec)
Parent     : $task.Wait($TimeoutMilliSec)

Arguments  : 
Expression : $client
Member     : Close
Static     : False
StaticType : System.Object
Extent     : $client.Close()
Parent     : $client.Close()

Arguments  : 
Expression : $client
Member     : Dispose
Static     : False
StaticType : System.Object
Extent     : $client.Dispose()
Parent     : $client.Dispose()

As you can see, you now get extremely detailed contextual information: not only do you know the raw name of the .NET method, but you also know which variable provided the type and the arguments that were passed.

It would still take considerable effort to create a tool that traverses the AST upstream to eventually find the .NET type that provided the method, but the AST provides all the necessary information.

For example, in this case, the ConnectAsync method was provided by the $client variable . How could you ask the AST what was originally assigned to this variable?

Analyzing Variable Assignments

With the simple tokenizer, you could generate lists of variable names, but the tokenizer lacked context and couldn’t, for example, tell you the content of a variable. With the AST, this is now relatively simple, and you can just repeat the steps from above. This time, we’re interested in the AssignmentStatementAst.

$path = "C:\test\some_file.ps1"

# get AST for script file
$errors = 
$tokens = $null
$ast = [System.Management.Automation.Language.Parser]::ParseFile($Path, [ref] $tokens, [ref]$errors)

# what would you like to know about this script?
$filter = { 
  param($astObject)  
  
  $astObject -is [System.Management.Automation.Language.AssignmentStatementAst] 
}

# search for all AST objects, including nested script blocks
$recurse = $true
   
# find the AST objects
$ast.FindAll($filter, $recurse)  

The result is a list of all assignments performed in the script you’re examining:

Left          : $client
Operator      : Equals
Right         : [System.Net.Sockets.TcpClient]::new()
ErrorPosition : =
Extent        : $client = [System.Net.Sockets.TcpClient]::new()
Parent        : {
                            $client = [System.Net.Sockets.TcpClient]::new()
                            $task = $client.ConnectAsync($ComputerName, $Port)
                            if ($task.Wait($TimeoutMilliSec))
                            {
                                $success = $client.Connected
                            }
                            else
                            {
                                $success = $false
                            }
                        }

Left          : $task
Operator      : Equals
Right         : $client.ConnectAsync($ComputerName, $Port)
ErrorPosition : =
Extent        : $task = $client.ConnectAsync($ComputerName, $Port)
Parent        : {
                            $client = [System.Net.Sockets.TcpClient]::new()
                            $task = $client.ConnectAsync($ComputerName, $Port)
                            if ($task.Wait($TimeoutMilliSec))
                            {
                                $success = $client.Connected
                            }
                            else
                            {
                                $success = $false
                            }
                        }

Left          : $success
Operator      : Equals
Right         : $client.Connected
ErrorPosition : =
Extent        : $success = $client.Connected
Parent        : {
                                $success = $client.Connected
                            }

Left          : $success
Operator      : Equals
Right         : $false
ErrorPosition : =
Extent        : $success = $false
Parent        : {
                                $success = $false
                            }  

This provides you with a clear list:

$ast.FindAll($filter, $recurse) | Select-Object -Property Left, Right

Here is the result for the example script I examined:

Left     Right                                     
----     -----                                     
$client  [System.Net.Sockets.TcpClient]::new()     
$task    $client.ConnectAsync($ComputerName, $Port)
$success $client.Connected                         
$success $false                                    
$success $false     

If you only wanted to know the assignment to $client, you could have asked the AST directly using its flexible script block filter:

$ast.Find( {param($x) $x -is [System.Management.Automation.Language.AssignmentStatementAst] -and $x.Operator -eq 'Equals' -and $x.Left.VariablePath.UserPath -eq 'client'  }, $true).Right    

Here is the result:

Expression                            Redirections Extent                                Parent                                         
----------                            ------------ ------                                ------                                         
[System.Net.Sockets.TcpClient]::new() {}           [System.Net.Sockets.TcpClient]::new() $client = [System.Net.Sockets.TcpClient]::new() 

This is how you can ask the AST for the content of the $client variable:

$ast.Find( {param($x) $x -is [System.Management.Automation.Language.AssignmentStatementAst] -and $x.Operator -eq 'Equals' -and $x.Left.VariablePath.UserPath -eq 'client'  }, $true).Right.Expression.Extent.Text  

Of course, this is somewhat convoluted now. Play around with the sample code and use your editor's IntelliSense to get a feel for the objects.
Here’s the cleaner code that retrieves all the variable assignments from a script:

$path = "C:\test\some_file.ps1"


# get AST for script file
$errors = $null
$tokens = $null
$ast = [System.Management.Automation.Language.Parser]::ParseFile($Path, [ref] $tokens, [ref]$errors)

# get all variable assignments
$filter = { 
  param($astObject)  
  
  $astObject -is [System.Management.Automation.Language.AssignmentStatementAst] -and
  $astObject.Operator -eq 'Equals'
}

# find the assignments
$ast.FindAll($filter, $true) |
  ForEach-Object {
    [PSCustomObject]@{
      Line = $_.Right.Expression.Extent.StartLineNumber
      Column = $_.Right.Expression.Extent.StartColumnNumber
      Variable = $_.Left.VariablePath.UserPath
      Assignment = $_.Right.Expression.Extent.Text      
    }
  }  

Related links 

• ScriptRunners Script collection provides ready-to-use PowerShell scripts.
• Try out ScriptRunner here
• ScriptRunner: Book a demo with our product experts