<#
.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, defaults to Ansible templated variable 'xylok_client_id'
    [Parameter(Mandatory = $false)]
    [String]$ClientID = '{{ xylok_client_id }}',
    
    # 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).ToUpper())
)

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
}

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
}

$DomainFQDN = Get-WmiObject -Query 'select domain from win32_computersystem' | Select-Object -ExpandProperty domain
$BaseURL = "https://xylok.$($domainFQDN.tolower())"
$BasePath = "$env:systemdrive\Xylok"
$RestHeaders = @{ 'Authorization' = "Bearer $Token" }

Start-Transcript -Path $basePath\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 and should be uppercase in Xylok."
    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 = $BasePath + '\' + $Filename
try {
    if ( !(Test-Path $BasePath) ) { New-Item -Path $BasePath -ItemType Directory -Force | Out-Null }
    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 = $BasePath + '\' + [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 {
    $checkScriptJob = Start-Job -ScriptBlock { powershell -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

$fileHeader = New-Object System.Net.Http.Headers.ContentDispositionHeaderValue 'form-data'
$fileHeader.Name = "`"filename`""
$fileHeader.FileName = "`"$($resultsFile.Name)`""

try {
    $packageFileStream = New-Object System.IO.FileStream @($resultsFile.FullName, [System.IO.FileMode]::Open)
    $streamContent = New-Object System.Net.Http.StreamContent $packageFileStream
    $streamContent.Headers.ContentDisposition = $fileHeader
    $streamContent.Headers.ContentType = New-Object System.Net.Http.Headers.MediaTypeHeaderValue 'application/octet'
}
catch {
    Write-Error 'Could not open/access results file.'
    Stop-Transcript
    return
}
   
$boundary = ('BoundaryPSXylok_' + [guid]::NewGuid().tostring())
$content = New-Object System.Net.Http.MultipartFormDataContent $boundary
$content.Add($streamContent)
$body = $content.ReadAsStringAsync().Result
$packageFileStream.Close()

$TaskID = Invoke-RestMethod -Uri "$BaseURL/api/v1/upload/" -Method Post -Headers $RestHeaders -Body $body -ContentType "multipart/form-data; boundary=$boundary"
$TaskID = $TaskID[0]
$Result = Watch-XylokTask -TaskID $TaskID -TaskDescription $TaskDescription
$ScanAddr = $Result.redirect
$ScanID = $ScanAddr.Trim('/').Split('/')[-1]

# request copying of existing answers if available
Write-Host 'Copying Previous Answers' -NoNewline
$PreviousScans = Invoke-RestMethod -Uri "$BaseURL/api/v1/machine/$MachineID/scans/?limit=1&offset=1" -Headers $RestHeaders -Method Get
if ( $PreviousScans.count -le 1 ) {
    Write-Host @skipped
    Write-Host @eolDone
}
else {
    $PreviousScanID = $PreviousScans.results[0].pk
    try {
        Invoke-RestMethod -Uri "$BaseURL/api/v1/scans/$PreviousScanID/copy-answers-to/$ScanID/" -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/$ScanID/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