Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Export OMS logs from an updated workspace (using Kusto language and API)
#Requires -Version 5.0
#Requires -Modules AzureRM.Resources, AzureRM.Profile, AzureServicePrincipalAccount
<#
=================================================================================
AUTHOR: Tao Yang
DATE: 14/11/2017
Version: 0.1
Comment: Export OMS logs from an updated workspace (using Kusto language and API)
=================================================================================
#>
[CmdletBinding()]
Param (
[Parameter(Mandatory = $true)][PSCredential]$AzureCredential,
[Parameter(Mandatory = $true)]
[ValidateScript({
try {
[System.Guid]::Parse($_) | Out-Null
$true
} catch {
$false
}
})]
[String]$TenantId,
[Parameter(Mandatory = $true)]
[ValidateScript({
try {
[System.Guid]::Parse($_) | Out-Null
$true
} catch {
$false
}
})]
[string]$WorkspaceId,
[Parameter(Mandatory = $true)][ValidateNotNullOrEmpty()][string]$SearchQuery,
[Parameter(Mandatory = $false)][ValidateNotNullOrEmpty()][DateTime]$StartUTCTime,
[Parameter(Mandatory = $false)][ValidateNotNullOrEmpty()][DateTime]$EndUTCTime,
[Parameter(Mandatory = $false)][Validaterange(1,600)][int]$Timeout = 180,
[Parameter(Mandatory = $false)][ValidateScript({Test-Path $_})][string]$OutputDir = $PSScriptRoot,
[Parameter(Mandatory = $false)][ValidateNotNullOrEmpty()][string]$OutputFileNamePrefix = 'OMSSearchResult',
[Parameter(Mandatory = $false)][ValidateSet('JSON', 'CSV')][string]$OutputFormat = 'CSV',
[Parameter(Mandatory = $false)][Validaterange(1000,10000)][int]$MaximumRowPerFile = 5000
)
#region functions
Function ConvertFrom-LogAnalyticsJson
{
[CmdletBinding()]
[OutputType([Object])]
Param (
[parameter(Mandatory=$true)]
[string]$JSON
)
$data = ConvertFrom-Json $JSON
$count = 0
foreach ($table in $data.Tables) {
$count += $table.Rows.Count
}
$objectView = New-Object object[] $count
$i = 0;
foreach ($table in $data.Tables) {
foreach ($row in $table.Rows) {
# Create a dictionary of properties
$properties = @{}
for ($columnNum=0; $columnNum -lt $table.Columns.Count; $columnNum++) {
$properties[$table.Columns[$columnNum].name] = $row[$columnNum]
}
# Then create a PSObject from it. This seems to be *much* faster than using Add-Member
$objectView[$i] = (New-Object PSObject -Property $properties)
$null = $i++
}
}
$objectView
}
Function Invoke-OMSKustoSearch
{
Param (
[Parameter(Mandatory = $true)][string]$AADToken,
[Parameter(Mandatory = $true)]
[ValidateScript({
try {
[System.Guid]::Parse($_) | Out-Null
$true
} catch {
$false
}
})]
[string]$WorkspaceId,
[Parameter(Mandatory = $true)][ValidateNotNullOrEmpty()][string]$SearchQuery,
[Parameter(Mandatory = $false)][Validaterange(1,600)][int]$Timeout = 180,
[Parameter(Mandatory = $false)][ValidateNotNullOrEmpty()][string]$ISO8601TimeSpan
)
#Constructing queries
$arrSearchResults = @()
$RowNumber = 0
$InitialQuery = "$SearchQuery | sort by TimeGenerated asc | extend rn=row_number()"
$subsequentQueryTemplate = "$InitialQuery | where rn > {0}"
$OMSAPIResourceURI = 'https://api.loganalytics.io'
$OMSAPISearchURI = "$OMSAPIResourceURI/v1/workspaces/$WorkspaceId/query"
#request header
$RequestHeader = @{
'Authorization' = $AADToken
"Content-Type" = 'application/json'
'prefer' = "wait=$Timeout, v1-response=true"
}
#intial query request
Write-Verbose "invoking intial search request using query `"$InitialQuery`""
Write-Verbose "Timespan: '$ISO8601TimeSpan'"
#Construct REST request body
$RequestBody = @{
"query" = $InitialQuery
}
If ($PSBoundParameters.ContainsKey('ISO8601TimeSpan'))
{
$RequestBody.Add('timespan', $ISO8601TimeSpan)
}
$RequestBodyJSON = ConvertTo-Json -InputObject $RequestBody
#Invoke search REST request
$SearchRequest = Invoke-WebRequest -UseBasicParsing -Uri $OMSAPISearchURI -Headers $RequestHeader -Body $RequestBodyJSON -Method Post -Verbose
#process result
Write-Verbose "Parsing Log Analytics Query REST API Results."
$arrSearchResults += ConvertFrom-LogAnalyticsJson $SearchRequest.Content
Write-Verbose "Number of rows retrieved so far: $($arrSearchResults.count)"
#Check if subsequent requests are required
$objResponse = ConvertFrom-JSON $SearchRequest.Content
If ($objResponse.error -ne $null)
{
Write-Verbose 'Initial query did not complete successful. Potentially hitting the API throttling limits.'
Write-Verbose " - Error Code: $($objResponse.error.code)"
Write-Verbose " - Error Message: $($objResponse.error.message)"
Write-Verbose " - Inner Error code: $($objResponse.error.details.innererror.code)"
Write-Verbose " - Inner Error Message: $($objResponse.error.details.innererror.message)"
$iRepeat = 0
If ($objresponse.error.code -ieq 'partialerror')
{
$iRepeat ++
Write-Verbose "Partial Error occurred, subsequent queries required. Repeat count: $iRepeat."
$bQueryCompleted = $false
Do
{
Write-Verbose "Getting the row number for the last row returned from all previous requests."
$RowNumber = $RowNumber + $arrSearchResults.count
$subsequentQuery = [string]::Format($subsequentQueryTemplate, $RowNumber)
Write-Verbose "Performing subsequent query number $iRepeat using query `"$subsequentQuery`""
$SubsequentRequestBody = @{
"query" = $subsequentQuery
}
If ($PSBoundParameters.ContainsKey('ISO8601TimeSpan'))
{
$SubsequentRequestBody.Add('timespan', $ISO8601TimeSpan)
}
$SubsequentRequestBodyJSON = ConvertTo-Json -InputObject $SubsequentRequestBody
#Invoke search REST request
$SubsequentSearchRequest = Invoke-WebRequest -UseBasicParsing -Uri $OMSAPISearchURI -Headers $RequestHeader -Body $SubsequentRequestBodyJSON -Method Post
#process result
Write-Verbose "Parsing Log Analytics Query REST API Results."
$arrSearchResults += ConvertFrom-LogAnalyticsJson $SubsequentSearchRequest.Content
#Check if subsequent requests are required
$objSubsequentResponse = ConvertFrom-JSON $SubsequentSearchRequest.content
if ($objSubsequentResponse.error.code -ine 'partialerror') {$bQueryCompleted = $true}
} while (!$bQueryCompleted)
Write-Verbose "Subsequent queries completed successful."
}
} else {
Write-Verbose "Initial search query retrieved the entire result set."
}
$arrSearchResults
}
Function Get-QueryTimeSpan
{
[CmdletBinding()]
Param (
[Parameter(Mandatory = $false)][ValidateNotNullOrEmpty()][DateTime]$StartUTCTime,
[Parameter(Mandatory = $false)][ValidateNotNullOrEmpty()][DateTime]$EndUTCTime
)
$UTCNow = (get-date).ToUniversalTime()
If (!$PSBoundParameters.ContainsKey('StartUTCTime'))
{
#If start time not specified, use 1 day as it is the default period if searching from OMS portal
$StartUTCTime = $UTCNow.AddDays(-1)
Write-Verbose "Start UTC Time noe specified. using the default value which is 1 day ago: '$StartUTCTime'"
}
If (!$PSBoundParameters.ContainsKey('EndUTCTime'))
{
#If end time not specified, use the current UTC time
$EndUTCTime = $UTCNow
Write-Verbose "Start UTC Time noe specified. using the current date time: '$EndUTCTime'"
}
$StartYear = $StartUTCTime.Year
$StartMonth = '{0:D2}' -f $StartUTCTime.Month
$StartDay = '{0:D2}' -f $StartUTCTime.Day
$StartHour = '{0:D2}' -f $StartUTCTime.Hour
$StartMinute = '{0:D2}' -f $StartUTCTime.Minute
$StartSecond = '{0:D2}' -f $StartUTCTime.Second
$EndYear = $EndUTCTime.Year
$EndMonth = '{0:D2}' -f $EndUTCTime.Month
$EndDay = '{0:D2}' -f $EndUTCTime.Day
$EndHour = '{0:D2}' -f $EndUTCTime.Hour
$EndMinute = '{0:D2}' -f $EndUTCTime.Minute
$EndSecond = '{0:D2}' -f $EndUTCTime.Second
$ISO8601TimeSpanTemplate = "{0}-{1}-{2}T{3}:{4}:{5}Z/{6}-{7}-{8}T{9}:{10}:{11}Z"
$ISO8601TimeSpan = [System.String]::Format($ISO8601TimeSpanTemplate, $StartYear, $StartMonth, $StartDay, $StartHour, $StartMinute, $StartSecond, $EndYear, $EndMonth, $EndDay, $EndHour, $EndMinute, $EndSecond)
$ISO8601TimeSpan
}
Function Export-ResultToFile
{
[CmdletBinding()]
Param (
[Parameter(Mandatory = $true)][psobject[]]$Logs,
[Parameter(Mandatory = $true)][ValidateScript({Test-Path $_})][string]$OutputDir,
[Parameter(Mandatory = $true)][ValidateNotNullOrEmpty()][string]$FileName,
[Parameter(Mandatory = $true)][ValidateSet('JSON', 'CSV')][string]$OutputFormat
)
$OutputFilePath = Join-Path $OutputDir $FileName
Write-Verbose "Exporting to '$OutputFilePath'..."
Switch ($OutputFormat)
{
'CSV' {$Logs | Export-CSV -LiteralPath $OutputFilePath -NoTypeInformation -Force}
'JSON' {ConvertTo-JSON -InputObject $Logs | Out-File $OutputFilePath -Force}
}
$OutputFilePath
}
#endregion
#region variables
$OMSAPIResourceURI = 'https://api.loganalytics.io'
#endregion
#region main
#Get AAD Token
Write-Verbose "Requesting Azure AD oAuth token"
$AADToken = AzureServicePrincipalAccount\Get-AzureADToken -TenantID $TenantId -Credential $AzureCredential -ResourceURI $OMSAPIResourceURI
#Work out the search time span
$TimeSpanParam = @{}
If ($PSBoundParameters.ContainsKey('StartUTCTime'))
{
$TimeSpanParam.Add('StartUTCTime', $StartUTCTime)
}
If ($PSBoundParameters.ContainsKey('EndUTCTime'))
{
$TimeSpanParam.Add('EndUTCTime', $EndUTCTime)
}
$ISO8601TimeSpan = Get-QueryTimeSpan @TimeSpanParam
Write-Output "Log Analytics search request ISO 8601 time span: '$ISO8601TimeSpan'."
#Invoke search API
Write-Verbose "Invoking search request. Search query: `"$SearchQuery`"... This could take a while"
$SearchResult = Invoke-OMSKustoSearch -AADToken $AADToken -WorkspaceId $WorkspaceId -SearchQuery $SearchQuery -ISO8601TimeSpan $ISO8601TimeSpan -Timeout $Timeout
Write-Output "Total number of rows returned: $($SearchResult.count)"
#Export logs to files
$bAllExported = $false
$iExportCount = 0
$totalExported = 0
$arrExportFiles = @()
Do {
$ExportSet = $searchresult | Select-Object -First $MaximumRowPerFile -Skip ($MaximumRowPerFile * $iExportCount)
$iExportCount ++
#Export Log
$FileName = "$OutputFileNamePrefix-$($ISO8601TimeSpan.replace(':', '.').split('Z/')[0])`-$iExportCount`.$OutputFormat"
Write-output "Exporting Batch No. $iExportCount with $($ExportSet.count) rows to '$FileName'"
$arrExportFiles += Export-ResultToFile -Logs $ExportSet -OutputDir $OutputDir -FileName $FileName -OutputFormat $OutputFormat -Verbose
$totalExported = $totalExported + $ExportSet.Count
If ($totalExported -eq $SearchResult.Count)
{
$bAllExported = $true
}
} While (!$bAllExported)
Write-output '', "All logs are exported to '$OutputDir'."
Write-Output "Total files created: $($arrExportFiles.Count)"
Write-output "Done!"
#endregion
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment