<#
.SYNOPSIS
    Collects the data protection state of a Windows computer.

.DESCRIPTION
    This script creates an archive of various system information that describes the data
    protection (DP) state at a particular point in time. It's intended to be run before and after
    reproducing an issue where a loss of protected data has occurred (e.g. the loss of Beyond
    Identity keys after a password reset).

.PARAMETER OutputDir
    Specifies the directory where the archive will be saved. The default is the same directory as
    the script.

.EXAMPLE
    .\CollectDPState.ps1

.OUTPUTS
    A timestamped zip archive containing collected data in a JSON file, and the copied DP files.

.NOTES
    Data collected:
      - Basic computer and operating system information
      - TPM information, including the SRK and EK public keys
      - DPAPI provider registry values
      - DPAPI master key files (the contents of %APPDATA%\Microsoft\Protect)
      - Domain controller information
      - Applications of interest
#>
param(
    [string]$OutputDir = $PSScriptRoot
)

class ComputerSystemInfo {
    [string]$Domain
    [string]$Manufacturer
    [string]$Model
    [bool]$PartOfDomain
    [string]$SystemType
    [string]$UserName
}

class OperatingSystemInfo {
    [string]$BuildNumber
    [string]$Caption
    [string]$LastBootUpTime
    [string]$OSArchitecture
    [uint16]$ServicePackMajorVersion
    [uint16]$ServicePackMinorVersion
    [string]$Version
}

class TpmInfo {
    [string]$Version
    [string]$Level
    [string]$Revision
    [string]$VendorId
    [string]$Firmware
    [bool]$IsBitLockerEnabled
    [string]$SrkPub
    [string]$EkPub
}

class RegistryKeyData {
    [string]$KeyName
    [System.Collections.Specialized.OrderedDictionary]$Values
}

class DomainController {
    [string]$Hostname
    [string]$IpV4Address
    [string]$OS
    [string]$Site
    [string]$Usn
    [bool]$IsRpcPortOpen
    [bool]$IsLdapPortOpen
    [bool]$IsSmbPortOpen
    [bool]$IsGCPort1Open
    [bool]$IsGCPort2Open
}

class DomainState {
    [string]$ForestName
    [int]$DomainControllerCount
    [System.Collections.Generic.List[DomainController]]$DomainControllers
}

class Application {
    [string]$Name
    [string]$Version
}

class ApplicationState {
    [System.Collections.Generic.List[Application]]$Applications
}

class DataProtectionState {
    [ComputerSystemInfo]$ComputerSystemInfo
    [OperatingSystemInfo]$OperatingSystemInfo
    [TpmInfo]$TpmInfo
    [System.Collections.Generic.List[RegistryKeyData]]$RegistryKeys
    [DomainState]$DomainState
    [ApplicationState]$ApplicationState
}

function GetComputerSystemInfo {
    $cs = Get-WmiObject -Class Win32_ComputerSystem

    $info = [ComputerSystemInfo]::new()
    $info.Domain = $cs.Domain
    $info.Manufacturer = $cs.Manufacturer
    $info.Model = $cs.Model
    $info.PartOfDomain = $cs.PartOfDomain
    $info.SystemType = $cs.SystemType
    $info.UserName = $cs.UserName

    $info
}

function GetOperatingSystemInfo {
    $os = Get-WmiObject -Class Win32_OperatingSystem

    $info = [OperatingSystemInfo]::new()
    $info.BuildNumber = $os.BuildNumber
    $info.Caption = $os.Caption
    $info.LastBootUpTime = $os.LastBootUpTime
    $info.OSArchitecture = $os.OSArchitecture
    $info.ServicePackMajorVersion = $os.ServicePackMajorVersion
    $info.ServicePackMinorVersion = $os.ServicePackMinorVersion
    $info.Version = $os.Version

    $info
}

function IsBitLockerEnabled {
    $shell = New-Object -ComObject Shell.Application
    $status = $shell.Namespace("C:").Self.ExtendedProperty("System.Volume.BitLockerProtection")
    $status -eq 1
}

function BufferToString ([byte[]]$buffer) {
    [System.Text.Encoding]::Unicode.GetString($buffer)
}

function BufferToBase64String ([byte[]]$buffer) {
    [Convert]::ToBase64String($buffer)
}

function IntToHex ([int]$n) {
    '0x{0:X}' -f $n
}

function GetTpmInfo {
    $ncryptDef = @'
    [DllImport("ncrypt.dll", SetLastError = true)]
    public static extern int NCryptOpenStorageProvider(
        ref IntPtr phProvider,
        [MarshalAs(UnmanagedType.LPWStr)] string pszProviderName,
        int dwFlags
    );

    [DllImport("ncrypt.dll", SetLastError = true)]
    public static extern int NCryptGetProperty(
        IntPtr hObject,
        [MarshalAs(UnmanagedType.LPWStr)] string pszProperty,
        [MarshalAs(UnmanagedType.LPArray)] byte[] pbOutput,
        int cbOutput,
        ref int pcbResult,
        int dwFlags
    );

    [DllImport("ncrypt.dll", SetLastError = true)]
    public static extern int NCryptFreeObject(
        IntPtr hObject
    );
'@

    $NCrypt = Add-Type -MemberDefinition $ncryptDef -Namespace Win32 -Name NCrypt -PassThru

    function Failed ([int]$status) {
        $status -lt 0
    }

    function NCryptGetProperty ([IntPtr]$prov, [string]$prop) {
        $dwSize = 0
        $status = $NCrypt::NCryptGetProperty($prov, $prop, $null, 0, [ref]$dwSize, 0)
        if (Failed $status) {
            throw "NCryptGetProperty failure (property: $prop, error: $(IntToHex $status))"
        }

        $buffer = [byte[]]::new($dwSize)
        $status = $NCrypt::NCryptGetProperty($prov, $prop, $buffer, $buffer.Length, [ref]$dwSize, 0)
        if (Failed $status) {
            throw "NCryptGetProperty failure (property: $prop, error: $(IntToHex $status))"
        }

        $buffer
    }

    $tpmInfo = [TpmInfo]::new()
    $tpmInfo.IsBitLockerEnabled = IsBitLockerEnabled

    try {
        $provName = "Microsoft Platform Crypto Provider"
        [IntPtr]$prov = [IntPtr]::Zero
        $status = $NCrypt::NCryptOpenStorageProvider([ref]$prov, $provName, 0)
        if (Failed $status) {
            throw "NCryptOpenStorageProvider failure (provider: $provName, error: $(IntToHex $status))"
        }

        # Returns a string in the form of:
        # TPM-Version:2.0 -Level:0-Revision:1.38-VendorID:'STM '-Firmware:65793.0
        $tpmStr = BufferToString (NCryptGetProperty $prov "PCP_PLATFORM_TYPE")
        $vals = $tpmStr.Split('-')
        $tpmInfo.Version  = ($vals[1].Split(':'))[1]
        $tpmInfo.Level    = ($vals[2].Split(':'))[1]
        $tpmInfo.Revision = ($vals[3].Split(':'))[1]
        $tpmInfo.VendorId = ($vals[4].Split(':'))[1]
        $tpmInfo.Firmware = ($vals[5].Split(':'))[1]

        $tpmInfo.SrkPub = BufferToBase64String (NCryptGetProperty $prov "PCP_SRKPUB")
        $tpmInfo.EkPub = BufferToBase64String (NCryptGetProperty $prov "PCP_EKPUB")
    }
    catch {
        throw
    }
    finally {
        if ($prov -ne [IntPtr]::Zero) {
            [void]$NCrypt::NCryptFreeObject($prov)
        }
    }

    $tpmInfo
}

function GetRegistryKey ([string]$path) {
    try {
        $key = Get-Item -Path $path -ErrorAction Stop
    }
    catch {
        throw "Failed to get registry key (path: $path, error: $_)"
    }
    
    $values = [ordered]@{}
    foreach ($name in $key.GetValueNames() | Sort-Object) {
        $values.Add($name, $key.GetValue($name))
    }

    $keyData = [RegistryKeyData]::new()
    $keyData.KeyName = $key.Name
    $keyData.Values = $values

    $key.Dispose()
    $keyData
}

function GetDomainState {
    Write-Host "Getting domain state..."
	$cs = Get-WmiObject -Class Win32_ComputerSystem

    if ($cs.PartOfDomain -ne $true) {
        write-host -fore red "This computer is not domain joined."
        return
    }

    $domainState = [DomainState]::new()
    $domains = [System.Collections.Generic.List[DomainController]]::new()

	$context = new-object 'System.DirectoryServices.ActiveDirectory.DirectoryContext'("domain", $cs.Domain)
    $dcs = @([System.DirectoryServices.ActiveDirectory.DomainController]::FindAll($context))
    $domainState.DomainControllerCount = $dcs.Length
	foreach ($entry in $dcs) {
        $domainState.ForestName = $entry.Forest
        $dc = [DomainController]::new()
        $dc.IpV4Address = $entry.IPAddress
        $dc.Hostname = $entry.Name
        $dc.OS = $entry.OSVersion
        $dc.Usn = $entry.HighestCommittedUsn
        $dc.Site = $entry.SiteName
        Write-Host "Testing Domain Controller ports..."
        $dc.IsRpcPortOpen = (Test-NetConnection -ComputerName $dc.Hostname -Port 135).TcpTestSucceeded
        $dc.IsLdapPortOpen = (Test-NetConnection -ComputerName $dc.HostName -Port 389).TcpTestSucceeded
        $dc.IsSmbPortOpen = (Test-NetConnection -ComputerName $dc.HostName -Port 445).TcpTestSucceeded
        $dc.IsGCPort1Open = (Test-NetConnection -ComputerName $dc.HostName -Port 3268).TcpTestSucceeded
        $dc.IsGCPort2Open = (Test-NetConnection -ComputerName $dc.HostName -Port 3269).TcpTestSucceeded
        $domains.Add($dc)
    }
    $domainState.DomainControllers = $domains

    $domainState
}

function GetApplicationState ([string]$namePattern) {
    Write-Host "Getting application state..."
    $appState = [ApplicationState]::new()
    $appList = [System.Collections.Generic.List[Application]]::new()
    $apps = @(Get-WmiObject -Class Win32_Product | Where Name -Like $namePattern | Select Name, Version)
    foreach ($entry in $apps) {
        $app = [Application]::new()
        $app.Name = $entry.Name
        $app.Version = $entry.Version
        $appList.Add($app)
    }
    $appState.Applications = $appList

    $appState
}

# Returns a DataProtectionState object containing
# 1. Computer, OS, and TPM information, registry entries for the DPAPI provider.
# 2. Domain information
# 3. User information
# 4. Application information
function GetDataProtectionState {
    $dpState = [DataProtectionState]::new()
    $dpState.ComputerSystemInfo = GetComputerSystemInfo
    $dpState.OperatingSystemInfo = GetOperatingSystemInfo
    $dpState.TpmInfo = GetTpminfo
    $dpState.DomainState = GetDomainState
    $dpState.ApplicationState = GetApplicationState("*Sailpoint*")
    
    $regKeys = [System.Collections.Generic.List[RegistryKeyData]]::new()
    $regKeys.Add((GetRegistryKey 'HKLM:\SOFTWARE\Microsoft\Cryptography\Protect\Providers\df9d8cd0-1501-11d1-8c7a-00c04fc297eb'))
    $dpState.RegistryKeys = $regKeys

    $dpstate
}

# Copies the %APPDATA%\Microsoft\Protect directory and its contents to the destination directory.
function CopyDpFiles ([string]$destDir) {
    Copy-Item -Path "$env:APPDATA\Microsoft\Protect" -Destination $destDir -Recurse
}

# Returns a JSON string with better formatting than the default PS 5.1 ConvertTo-Json.
# Based on https://github.com/PowerShell/PowerShell/issues/2736.
function FormatJson ([string]$json) {
    $res = [System.Collections.Generic.List[string]]::new()
    $indent = 0
    foreach ($line in $json -Split [Environment]::NewLine) {
        if ($line -match '[\}\]]') {
            # Line contains } or ]
            $indent--
        }
        $fmtLine = (' ' * $indent * 4) + $line.TrimStart().Replace(':  ', ': ')
        if ($line -match '[\{\[]') {
            # Line contains { or [
            $indent++
        }
        $res.Add($fmtLine)
    }
    $res -Join [Environment]::NewLine
}

function CreateArchive ([string]$sourceDir, [string]$zipPath) {
    Add-Type -AssemblyName System.IO.Compression.FileSystem
    [System.IO.Compression.ZipFile]::CreateFromDirectory($sourceDir, $zipPath)
}

$ErrorActionPreference = 'Stop'
Write-Host 'Collecting DP state...'

$dpState = GetDataProtectionState

$timestamp = (Get-Date).ToUniversalTime().ToString("yyyyMMddTHHmmssffZ");
$tempDir = "$OutputDir\dp_state_temp_$timestamp"
New-Item -Path $tempDir -ItemType 'Directory' | Out-Null

CopyDpFiles $tempDir

$json = FormatJson ($dpState | ConvertTo-Json -Depth 3)
$json | Set-Content -Path "$tempDir\dp_state_$timestamp.json"

$zipPath = "$OutputDir\dp_state_$timestamp.zip"
CreateArchive $tempDir $zipPath
Remove-Item -Path $tempDir -Recurse -Force

Write-Host 'Data collection completed.' -ForegroundColor Green
Write-Host "Archive: $zipPath" -ForegroundColor Green