function Show-Usage {
	Write-Host "cbc-sensor-capture.ps1 captures performance data for CBC sensor"
	Write-Host "Script starts requested traces and pauses for use to reproduce specific issues or capture adequate sample."  
	Write-Host "Hit enter to stop trace and capture data."
	Write-Host "Needs to be run in account with admin privileges and access to authenticated CLI commands."
    Write-Host "If a low alt procmon is desired, first copy procmonlowalt.exe to c:\temp before running script."
	Write-Host " "
	Write-Host "Usage : cbc-sensor-capture.ps1 <command>"
	Write-Host "The following commands are currently supported. Note you can specify multiple commands except wpr and procmon in same run is not allowed."
	Write-Host "amsi           - enables AMSI ETL logging"
    Write-Host "bypass         - puts the sensor into bypass. useful for capturing procmon and wpr traces to compare with sensor enabled."
	Write-Host "config         - updates config settings stored in c:\temp\config.json. Note config changes will stay in place until sensor or system restarts."
	Write-Host "debug          - enables sensor verbos debuggging"
    Write-Host "keepevents     - retain psc event batches"
    Write-Host "netsh          - capture network packets" 
	Write-Host "procmon        - capture process monitor trace"
    Write-Host "skipreset      - counters are reset by default. If you do not want to reset countesr for test, specify this option"
	Write-Host "wpr            - capture Windows performance trace"
	Write-Host "minifilter     - capture Windows performance trace with minifilter I/O activity"
	Write-Host " "
	Write-Host "Script works better if run from regular command line (cmd.exe)"
}

function Test-Admin {
	$principal = [Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()
	if( -not $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
		Write-Host "You must be an Administrator to run this tool"
		exit 1
	}	
}

function Test-Repcli {
	& "C:\Program Files\Confer\RepCLI.exe" configprops | Out-Null
	if ( !$? )
	{
		Write-Host  "Not running as CLI authenticated user"
		exit 1
	}
}

function Reset-Counters {
	& "C:\Program Files\Confer\RepCLI.exe" resetcounters
	if ( !$? ) {
		Write-Host "Reset Sensor Counters failed"
	}
}

function Download-Procmon {
	[string]$Destination = "c:\temp\procmon64.exe"
	if (-not (Test-Path $Destination)) {
		(New-Object System.Net.WebClient).DownloadFile("https://live.sysinternals.com/procmon64.exe", "c:\temp\procmon64.exe")
	}
    if ( [System.IO.File]::Exists( "c:\temp\procmon64.exe" )) {
        Set-Variable -name "downloadProcmon" -Value $true -Scope global
        return "c:\temp\procmon64.exe"
    }
    else {
        return $null
    }
}

function Get-Procmon {
    if ( [System.IO.File]::Exists("c:\temp\procmonlowalt.exe") ) {
        return "c:\temp\procmonlowalt.exe"
    } elseif ( [System.IO.File]::Exists("c:\temp\procmon64.exe") ) {
        return "c:\temp\procmon64.exe"
    } elseif ( [System.IO.File]::Exists( "c:\temp\procmon.exe") ) {
        return "c:\temp\procmon.exe"
    }
    #$downloadProcmon = [System.Windows.MessageBox]::Show('Unable to find procmon in c:\temp. Would  you like to download procmon from sysinternsl?','Game input','YesNo','Error')
    #$cont = Read-Host -Prompt "Procmon not found in c:\temp. Do you want to download from live.sysinternals.com?"
    #if ( $downloadpProcmon -eq 'yes') {
    $cont = Read-Host -Prompt "Procmon not found in c:\temp. Do you want to download from live.sysinternals.com?"
    if ( $cont -eq 'yes' -or $cont -eq 'y' ) {
        return Download-Procmon
    }
    else {
        return $null
    }
}

# main

# Policies to enable/disable event batch retention
$keepEventsOn = @"
{ 
 "Version": "1.0",
 "ConfigProps": [
  {
   "Name": "PscEventBatchKeepLocalCopy",
   "Data": "1"
  }
 ]
}
"@

$keepEventsOff = @"
{ 
 "Version": "1.0",
 "ConfigProps": [
  {
   "Name": "PscEventBatchKeepLocalCopy",
   "Data": "0"
  }
 ]
}
"@

# globals
$global:downloadProcmon = $null

$resetcounters = $false
$procmon = $false
$wpr = $false
$minifilter = $false
$debug = $false 
$procmonFile  = $null
$bypass = $false
$netsh = $false
$keepEvents = $false
$config = $false
$amsi = $false
$timeStamp=Get-Date -Format "yyyyddMMHHmmss"
$logFolder="c:\temp\"+$timeStamp
$pmlFile=$logFolder+"\sensor.pml"
$nettraceEtl=$logFolder+"\NetTrace.etl"
$nettraceMan=$logFolder+"\NetTrace.man"
$wprFile=$logFolder+"\sensor.etl"
$amsiFile=$logFolder+"\cbamsi.etl"
$capi2File=$logFolder+"\capi2.evtx"
$keepEventsJson=$logFolder+"\pscEventBatch.json"
$hostName=[System.Net.DNS]::GetHostByName(($env:computerName)).HostName

$uninstallCode = $null
$needUnlock = $false
$configJsonFile = "c:\temp\config.json"

try {
	$repmgrVersion = (get-item "c:\program files\confer\repmgr.exe").VersionInfo.FileVersion
	[float]$repmgrMajorMinor = $repmgrVersion.substring(0,3)
}
catch {
	[float]$repmgrMajorMinor = 0
}

$validArgs=@("help","skipreset","procmon","wpr","minifilter","debug","bypass","netsh","keepevents","amsi","config")

if ( $args.count -eq 0 ) {
	$confirmation = Read-Host "No options means only counters will be reset. Do you want to continue (y/n)?"
	if ( $confirmation -ne 'y' ) {
		Write-Host "Type 'cbc-sensor-capture help' for usage"
		exit 0
	}
} else {
	for ( $i = 0; $i -lt $args.count; $i++ ) {
		if ( -not $validArgs.Contains($args[$i]) ) {
			Write-Host "Invalid option " $args[$i]
			Write-Host "Type 'cbc-sensor-capture help' for usage"
			exit 1;
		}		
		switch ($args[$i]) {
			"help" { Show-Usage; exit 0 }
			"skipreset" { $skipreset = $true; }
			"procmon" { $procmon = $true; }
			"wpr" { $wpr = $true; }
			"minifilter" { $minifilter = $true; }
			"debug" { $debug = $true; }
			"bypass" { $bypass = $true; }
			"netsh" { $netsh = $true; }
			"keepevents" { $keepEvents = $true; }
			"amsi" { $amsi = $true; }
			"config" { $config = $true; }
		}
	}
}

if (-not [System.Environment]::Is64BitOperatingSystem ) {
    Write-Host "Script only supports 64-bit operating systems."
    exit 1
}

Test-Admin
Test-Repcli	

# Check if OS is Elam supported
$ELAMSupportedBuild = 15063
$ELAMSupported = $false
$osver = (Get-ItemProperty -Path "HKLM:\SOFTWARE\Microsoft\Windows NT\CurrentVersion").CurrentBuildNumber
# Convert to integer
$osver = 0 + $osver
if ( $osver -ge $ELAMSupportedBuild ) {
	Write-Host "ELAM Supported"
	$ELAMSupported = $true
} else {
	Write-Host "ELAM Not Supported"
}


if ( $procmon -and $wpr ) {
    Write-Host "Unable to run procmon and wpr at same time"
    exit 1
}

if ( ($wpr -or $minifilter) -And -Not [System.IO.File]::Exists("c:\windows\system32\wpr.exe")) {
    Write-Host "C:\Windows\System32\wpr.exe not found"
    exit 1
}

if ( $wpr -and $minifilter ) {
	Write-Host "Select wpr or minifilter but not both"
	exit 1
}

if ( $amsi -and -not [System.IO.File]::Exists("c:\windows\system32\logman.exe" )) {
	Write-Host "C:\Windows\System32\logman.exe not found"
	exit 1
}

if ( -Not (Test-Path -Path "C:\temp") ) { 
    New-Item -Path "C:\" -Name "temp" -ItemType "directory" | Out-Null
}
New-Item -Path "C:\temp" -name $timeStamp -ItemType "directory" |  Out-Null

if ( $procmon ) {
	$procmonFile = Get-Procmon
    if ( !$procmonFile ) {
        Write-Host "Please download procmon and try again."
        exit 1
    }
}

if ($keepEvents -Or $config) {
	$needUnlock = $true
}

# Switching to Confer folder so we can directly call repcli.exe
Push-Location "C:\Program Files\Confer"

if ( $procmon -and $ELAMSupported ) {
	Write-Host "Unprotecting CbDefense"
	.\repcli.exe bypass 1
	.\repcli.exe registerProtectedSvcs 0
	Write-Host "Restarting CbDefense"
	.\repcli.exe stopCbServices
	Start-Sleep -Seconds 5
	Start-Service -Name "CbDefense"
	Start-Sleep -Seconds 5
	.\repcli.exe bypass 0
}

if ( $needUnlock ) {
	$uninstallCode = Read-Host "Enter the uninstall code to unlock restricted RepCLI commands."
	if ( !$uninstallCode ) {
		Write-Host "Uninstall code required."
		exit 1 
	}
	& "C:\Program Files\Confer\RepCLI.exe" unlock $uninstallCode *> $null
	if ( !$? ) {
        Write-Host "RepCLI unlock failed."
		exit 1
	}
}

if ( $keepEvents ) {
	Write-Host "Enabling PscEventBatchKeepLocalCopy."
	$keepEventsOn | Out-File $keepEventsJson -Encoding utf8
	.\repcli addpolicy psc $keepEventsJson
	if ( !$? )
	{
		Write-Host "RepCLI addpolicy to retain PSC events failed."
	}
}

if ( $config ) {
	if ( [System.IO.File]::Exists( $configJsonFile )) {	
		Write-Host "Update config settings from $configJsonFile"
		.\repcli addpolicy psc $configJsonFile
		if ( !$? ) {
			Write-Host "RepCLI addpolicy to set configs failed."
		}
	}
	else {
		Write-Host "File $configJsonFile not found."
		$config = $false;
	}
}

if ( $procmon ) {
    Write-Host "Disabling Tamper Protect policy"	
    .\repcli.exe deletepolicy 1DED7E47-CE4C-448E-AD01-6F4AC3CE7F5D
    if ( !$? )
    {
        Write-Host "RepCLI deletePolicy failed."
    }
}

if ( $keepEvents -or $config -or $procmon ) {
	Write-Host "Pausing 30 seconds to allow sensor to complete post-update policy activities."
	Start-Sleep -Seconds 30
}

if ( $wpr ) {
    Write-Host "Starting WPR capture"
    wpr -start GeneralProfile -start CPU -start Registry -start FileIO -start DiskIO -start Network -filemode
}

if ( $minifilter ) {
    Write-Host "Starting WPR capture"
    wpr -start GeneralProfile -start CPU -start Registry -start FileIO -start DiskIO -start Minifilter -start Network -filemode
}

if ( $netsh ) {
	Write-Host "Enabling CAPI2 events"
	$logName = 'Microsoft-Windows-CAPI2/Operational'
	$log = New-Object System.Diagnostics.Eventing.Reader.EventLogConfiguration $logName
	$saveCAPIState = $log.IsEnabled
	$saveCAPIMode = $log.LogMode
	$log.IsEnabled=$true
	$log.LogMode = 2 # 1 = AutoBackup 0 = Circular 2 = Retain
	$log.SaveChanges()	
    Write-Host "Starting netsh trace"
    netsh trace start scenario=InternetClient capture=yes traceFile=$nettraceEtl
}    

if ( $debug ) {
    Write-Host "Enabling debug"
    .\repcli.exe debug 1
    .\repcli.exe kerneltrace 4 -1
}

if ( $bypass ) {
    Write-Host "Sensor entering bypass mode"
    .\repcli bypass 1
}

if ( $amsi ) {
	logman start amsisession -p {aa0a544f-5938-41c1-b992-81a1a81d4fa2} -o $amsiFile -ets
}

if ( $procmon ) {
    Write-Host "Starting procmon"
    $procmonCmd = "$procmonFile /AcceptEula /NoFilter /Quiet /Minimized /BackingFile " + $pmlFile
    Invoke-Expression $procmonCmd
}

# Reset counters unless skipreset is specified
if ( $skipreset ) {
	Write-Host "Skipping counter reset"
} else {
    Write-Host "Resetting counters"
    .\repcli.exe resetcounters
}

if ( $repmgrMajorMinor -ge 3.8 ) {
	.\repcli comment "cbs-sensor-capature - pausing for reproduction"
}

Write-Host "Reproduce behavior. If a general performance issue, allow enough time to pass to capture sufficient samples but not long enough for procmon/wpr size to become unweildly."
Read-Host "Hit enter to stop trace and collect logs."
#$Shell = New-Object -ComObject "WScript.Shell"
#$Button = $Shell.Popup("Reproduce behavior. Hit OK stop trace and collect logs.", 0, "start-trace.ps1", 0)

if ( $repmgrMajorMinor -ge 3.8 ) {
	.\repcli comment "cbs-sensor-capature - ending reproduction"
}


# always capture counters even if resetcounters not specified
Write-Host "Capturing counters"
.\repcli.exe counters | Out-File $logFolder\counters.txt -Encoding utf8

if ( $debug ) {
    Write-Host "Turning off debug"
    .\repcli.exe debug 0
    .\repcli.exe kerneltrace 0
}

if ( $procmon ) {
    Write-Host "Stopping procmon"
    $procmonCmd = "$procmonFile /Terminate"
    Invoke-Expression $procmonCmd
	# wait for process to exit
	$procmonLeaf=[io.path]::GetFileNameWithoutExtension($procmonFile)
	$p = Get-Process $procmonLeaf -ea SilentlyContinue
	if ($p) {
		Wait-Process -Id $p.id
	}
	# Reset using paging file
	#$procmonCmd = "$procmonFile /AcceptEula /Quiet /Minimized /PagingFile /NoConnect"
	#Invoke-Expression $procmonCmd
	#$procmonCmd = "$procmonFile /WaitForIdle"
	#Invoke-Expression $procmonCmd
	#$procmonCmd = "$procmonFile /Terminate"
	#Invoke-Expression $procmonCmd
}

if ( $wpr -or $minifilter ) {
    Write-Host "Stopping WPR. This can take several minutes to complete."
    wpr -stop $wprFile
}

if ( $netsh ) {
	Write-Host "Exporting CAPI2 events"
	wevtutil epl Microsoft-Windows-CAPI2/Operational $capi2File
	Write-Host "Resetting CAPI2 event state"
	$log.IsEnabled = $saveCAPIState
	$log.LogMode = $saveCAPIMode
	$log.SaveChanges()
    Write-Host "Stopping netsh trace. This can take several minutes to complete."
    netsh trace stop
    tracerpt -l $nettraceEtl -export $nettraceMan
}

if ( $amsi ) {
	Write-Host "Capaturing AMSI ETL file"
	logman -stop -n amsisession -ets
}

Write-Host "Capturing sensor diagnostics"
.\repcli.exe capture $logFolder -noSysDump

if ( $procmon ) {
	if ( $ELAMSupported ) {
		Write-Host "Restoring CbDefense protected service status"
		.\repcli.exe bypass 1
		.\repcli.exe registerProtectedSvcs 1
		Write-Host "Restarting CbDefense"
		.\repcli.exe stopCbServices
		Start-Sleep -Seconds 5
		Start-Service -Name "CbDefense"
		Start-Sleep -Seconds 5
		.\repcli.exe bypass 0
	} else {
		Write-Host "Re-enabling tamper policy"
		.\repcli.exe restorepolicy
	}
}


if ($bypass -and !$procmon) {
    Write-Host "Disabling bypass mode."
    .\repcli bypass 0
}

# if procmon set, no need to disable PscEventBatchKeepLocalCopy
if ( $keepEvents -and $uninstallCode -and !$procmon ) {
	Write-Host "Disabling PscEventBatchKeepLocalCopy."
	$keepEventsOff | Out-File $keepEventsJson -Encoding utf8
	.\repcli addpolicy psc $keepEventsJson
}
	

try {
    Write-Host "Zipping files"
	Copy-Item $PSCommandPath -Destination $logFolder
    $zipFile="c:\temp\$hostName-$timeStamp.zip"
    $compress = @{
        Path = $logFolder
        CompressionLevel = "Fastest"
        DestinationPath = $zipFile
    }
    Compress-Archive @compress -ea stop
    Remove-Item -Recurse $logFolder -Force
    Write-Host "Please upload "$zipFile" to CB Vault"
}
catch {
    Write-Host "Zip failed, please zip $logFolder and upload to CB Vault"
}

if ( $procmon -And $downloadProcmon ) {
    Remove-Item $procmonFile
}

Pop-Location