<#
.SYNOPSIS
    Runs the Xylok checker script for the machine and uploads the results to Xylok
.DESCRIPTION
    Finds the machine in Xylok (if exists), downloads the check script, runs the check script,
    uploads the results, performs auto analysis, and prints the output if requested.
.EXAMPLE
    PS C:\> DownloadXylokScript.ps1 -Token <token>
    Runs the script without any extra output
    PS C:\> DownloadXylokScript.ps1 -Token <token> -Timeout 0 -MachineName mymachine
        -PrintScriptOutput
    Runs the script, invoking the checker script with no timeout (this is helpful for long-running
    checks), using a custom machine name other than this machines hostname, and printing the
    output of the check script to the console
.NOTES
    This particular script is formatted to work with Ansible as a templated file such that the
    variable 'xylok_client_id' can be used to prefill the client ID for finding the machine in Xylok

    Thanks to TL Skinner for supplying this example.
#>
param (
    # Xylok authentication token
    [Parameter(mandatory = $true)]
    [String]$Token,

    # Xylok client ID
    [Parameter(Mandatory = $true)]
    [String]$ClientID = '',

    # Xylok client ID, defaults to Ansible templated variable 'xylok_client_id'
    [Parameter(Mandatory = $true)]
    [String]$BaseURL = '',

    # If included, prints the output of the script to the console
    [Parameter(mandatory = $false)]
    [Switch]$PrintScriptOutput,

    # Number of seconds to wait for the check script to run, default is 3600 (60 minutes), 0
    # disables the timeout
    [Parameter(mandatory = $false)]
    [Int]$Timeout = 3600,

    # Specifies the machine name in Xylok to lookup and retrive check script, defaults to uppercase
    # hostname of the machine running this script
    [Parameter(mandatory = $false)]
    [String]$MachineName = ($env:COMPUTERNAME)
)

Add-Type -AssemblyName System.Net.Http
Add-Type -AssemblyName System.Web

$greenCheck = @{
    Object          = ' OK'
    ForegroundColor = 'Green'
    NoNewLine       = $true
}

$redX = @{
    Object          = ' FAIL'
    ForegroundColor = 'Red'
    NoNewLine       = $true
}

$skipped = @{
    Object          = ' SKIPPED'
    ForegroundColor = 'Gray'
    NoNewLine       = $true
}

$eolDone = @{
    Object = ' (done)'
}

$statusMarker = @{
    Object    = '.'
    NoNewLine = $true
}

[System.Net.ServicePointManager]::ServerCertificateValidationCallback = {$true} ;
Set-Alias -Name pwshexe -Value 'powershell'
if ((Get-Command "pwsh" -ErrorAction SilentlyContinue) -eq $null)
{
   Set-Alias -Name pwshexe -Value 'pwsh'
}


function Watch-XylokTask {
    param(
        [Parameter(mandatory = $true)]
        [String]$TaskID,

        [Parameter(mandatory = $true)]
        [String]$TaskDescription,

        [Parameter(mandatory = $false)]
        [String]$MessageKey = 'message',

        [Parameter(mandatory = $false)]
        [Switch]$ProceedOnFailure
    )
    Do {
        $result = Invoke-RestMethod -Uri ("$Script:BaseURL/api/v1/tasks/$TaskID/") -Method Get -Headers $Script:RestHeaders
        Start-Sleep -Seconds 1
        Write-Host @statusMarker
    } Until (($result.complete -eq $true) -or ($result.failed -eq $true))
    if ( $result.failed ) {
        Write-Host @redX
        Write-Host @eolDone
        Write-Error ('Msg: ' + $result.result.result_msg)
        if ( $ProceedOnFailure ) { break }
        else { exit 1 }
    }
    Write-Host @greenCheck
    Write-Host @eolDone
    return $result
}

$RestHeaders = @{ 'Authorization' = "Bearer $Token" }

Start-Transcript -Path checkScriptAutoRun.txt -Force

$machines = Invoke-RestMethod -Uri "$BaseURL/api/v1/machine/?name=$([System.Web.HttpUtility]::UrlEncode($MachineName))&client_id=$ClientID" -Method Get -Headers $restHeaders -ContentType 'application/json'

if ( $machines.count -eq 0 ) {
    Write-Host "Machine $MachineName not found in Xylok. Names are case-sensitive"
    return
}

if ( $machines.count -gt 1 ) {
    Write-Host "$($machines.count) machines found in Xylok for client ID $ClientID with same name ($MachineName). Names should be unique for machines in a client."
}

# grab the machine ID and store it
Write-Host 'Machine found in Xylok'
$MachineID = $machines.results[0].pk

# Get the check script bundle
$TaskDescription = 'Getting check script bundle'
Write-Host $TaskDescription -NoNewline
$TaskID = Invoke-RestMethod -Uri "$BaseURL/api/v1/machine/$MachineID/script/" -Headers $restHeaders
$Result = Watch-XylokTask -TaskID $TaskID -TaskDescription $TaskDescription
$RedirectAddr = $Result.download

Write-Host 'Extracting check script from bundle' -NoNewline
# convert the downloaded data to zip file
$Filename = $RedirectAddr.split('/')[-1]
$Filepath = 'script.zip'
try {
    Invoke-RestMethod -Method Get -Uri "$BaseURL$RedirectAddr" -Headers $RestHeaders -OutFile $Filepath
    Write-Debug "Saved script to '$Filepath'."
    Write-Host @greenCheck
    Write-Host @eolDone
}
catch {
    Write-Host @redX
    Write-Host @eolDone
    Write-Error 'Failed to save script file.'
    Stop-Transcript
    return
}

# unzip the script
if ( ![io.path]::GetExtension($Filepath) -eq 'zip' ) {
    Write-Error 'Unexpected file type.'
    return
}
$DestPath = [io.path]::GetFileNameWithoutExtension($Filepath)
Expand-Archive -Path $Filepath -DestinationPath $DestPath -Force

# delete other results files
try {
    Get-ChildItem -Path $DestPath\* -Include '*.xylok' | Remove-Item -Force -ErrorAction SilentlyContinue
}
catch {
    Stop-Transcript
}

# attempt to run the powershell script
Write-Host 'Running check script' -NoNewline
try {
    if ((Get-Command "pwsh" -ErrorAction SilentlyContinue) -eq $null) {
      $checkScriptJob = Start-Job -ScriptBlock { powershell -file "$($Using:DestPath)\xylok-collect.ps1" -ExecutionPolicy Unrestricted -NoLogo }
    } else {
      $checkScriptJob = Start-Job -ScriptBlock { pwsh -file "$($Using:DestPath)\xylok-collect.ps1" -ExecutionPolicy Unrestricted -NoLogo }
    }

    $count = $timeout
    while (($checkScriptJob.State -eq 'Running') -and ($checkScriptJob.State -ne 'NotStarted')) {
        Write-Host @statusMarker
        Start-Sleep 10
        if ( $timeout -eq 0 ) { continue } #unlimited timeout
        $count -= 10
        if ( $count -lt 0 ) { throw New-Object System.TimeoutException }
    }
}
catch [System.TimeoutException] {
    Stop-Job $checkScriptJob -Confirm:$false
    Write-Host @redX
    Write-Host @eolDone
    Write-Error 'Running script timed out'
    Stop-Transcript
    return
}
catch {
    Write-Host @redX
    Write-Host @eolDone
    Write-Error 'Failed to run the check script.'
    Stop-Transcript
    return
}
Write-Host @greenCheck
Write-Host @eolDone

# upload results back to server
$TaskDescription = 'Uploading Results'
Write-Host $TaskDescription -NoNewline

$resultsFile = Get-ChildItem $DestPath\* -Include *.xylok | Sort-Object LastWriteTime -Descending | Select-Object -First 1

try {
  Write-Host $resultsFile.FullName
  $body = Get-Content -Path $resultsFile.FullName -Encoding UTF8 -Raw
}
catch {
    Write-Error 'Could not open/access results file.'
    Stop-Transcript
    return
}

$TaskID = Invoke-RestMethod -Uri "$BaseURL/api/v3/uploads/result/raw" -Method Post -Headers $RestHeaders -Body $body
$TaskID = $TaskID.jobs[0].id
$Result = Watch-XylokTask -TaskID $TaskID -TaskDescription $TaskDescription

$Scans = Invoke-RestMethod -Uri "$BaseURL/api/v1/machine/$MachineID/scans/?limit=2&offset=0" -Headers $RestHeaders -Method Get
if ( $Scans.count -eq 0 ) {
    Write-Error 'Failed to find uploaded scan'
    Stop-Transcript
    return
}

$NewScanID = $Scans.results[0].pk

# request copying of existing answers if available
if ( $Scans.count -gt 1 ) {
  Write-Host 'Copying Previous Answers' -NoNewline
  $PreviousScanID = $Scans.results[1].pk
  try {
    Invoke-RestMethod -Uri "$BaseURL/api/v1/scans/$PreviousScanID/copy-answers-to/$NewScanID/" -Headers $RestHeaders -Method Post | Out-Null
    Write-Host @greenCheck
    Write-Host @eolDone
  }
  catch {
    Write-Host @redX
    Write-Host @eolDone
    Write-Error ('Msg: ' + $_.Message)
  }
}

# request auto analyze
$TaskDescription = 'Performing auto-analysis'
Write-Host 'Analyzing Results' -NoNewline
$TaskID = Invoke-RestMethod -Uri "$BaseURL/api/v1/scans/$NewScanID/aa/execute/" -Method Post -Headers $RestHeaders
Watch-XylokTask -TaskID $TaskID -TaskDescription $TaskDescription -ProceedOnFailure | Out-Null

# print output if requested
if ( $PrintScriptOutput ) {
  Write-Host 'Check Script Output:'
  Receive-Job $checkScriptJob
}

Stop-Transcript