# -*- coding: utf-8 -*-
from setuphelpers import *

r"""
Sources:
https://learn.microsoft.com/windows/security/operating-system-security/data-protection/bitlocker/faq
https://learn.microsoft.com/windows/security/operating-system-security/data-protection/bitlocker/bitlocker-basic-deployment
https://learn.microsoft.com/powershell/module/bitlocker/get-bitlockervolume

https://learn.microsoft.com/en-us/windows/win32/secprov/getconversionstatus-win32-encryptablevolume#parameters


BitLocker Powershell informations:
PS C:\waptdev\enable-bitlocker\windows> Get-BitLockerVolume | Get-Member -Name *

   TypeName: Microsoft.BitLocker.Structures.BitLockerVolume

Name                 MemberType Definition
----                 ---------- ----------
Equals               Method     bool Equals(System.Object obj)
GetHashCode          Method     int GetHashCode()
GetType              Method     type GetType()
ToString             Method     string ToString()
AutoUnlockEnabled    Property   System.Nullable[bool] AutoUnlockEnabled {get;}
AutoUnlockKeyStored  Property   System.Nullable[bool] AutoUnlockKeyStored {get;}
CapacityGB           Property   float CapacityGB {get;}
ComputerName         Property   string ComputerName {get;}
EncryptionMethod     Property   Microsoft.BitLocker.Structures.BitLockerVolumeEncryptionMethodOnGet EncryptionMethod {…
EncryptionPercentage Property   System.Nullable[float] EncryptionPercentage {get;}
KeyProtector         Property   System.Collections.ObjectModel.ReadOnlyCollection[Microsoft.BitLocker.Structures.BitLo…
LockStatus           Property   Microsoft.BitLocker.Structures.BitLockerVolumeLockStatus LockStatus {get;}
MetadataVersion      Property   int MetadataVersion {get;}
MountPoint           Property   string MountPoint {get;}
ProtectionStatus     Property   Microsoft.BitLocker.Structures.BitLockerVolumeProtectionStatus ProtectionStatus {get;}
VolumeStatus         Property   System.Nullable[Microsoft.BitLocker.Structures.BitLockerVolumeStatus] VolumeStatus {ge…
VolumeType           Property   Microsoft.BitLocker.Structures.BitLockerVolumeType VolumeType {get;}
WipePercentage       Property   System.Nullable[float] WipePercentage {get;}


"""
volumestatus_dict = {
    0: "FullyDecrypted",
    1: "FullyEncrypted",
    2: "EncryptionInProgress",
    3: "DecryptionInProgress",
    4: "EncryptionPaused",
    5: "DecryptionPaused",
}
# https://learn.microsoft.com/en-us/windows/win32/secprov/getencryptionmethod-win32-encryptablevolume#parameters
# EncryptionMethod: XtsAes128 = 6; XtsAes256 = 7 #
encryptionmethod_dict = {
    0: "None",
    1: "AES_128_WITH_DIFFUSER",
    2: "AES_256_WITH_DIFFUSER",
    3: "Aes128",
    4: "Aes256",
    5: "HARDWARE_ENCRYPTION",
    6: "XtsAes128",
    7: "XtsAes256",
    -1: "(uint32)",
}
target_encryption_method = 7
allow_swap_encryption_method = False  # Not implemented yet


def install():

    # Installing the package
    if get_powershell_str("Get-ComputerInfo", "BiosFirmwareType").upper() == "UEFI":
        # BiosFirmwareType: BIOS = 1; UEFI = 2
        print("OK: This computer BIOS boot in UEFI mode")
    else:
        error("ERROR: This computer BIOS does not boot in UEFI mode, you have to change BIOS mode to continue")

    tpm_state = check_tpm_state()
    bitlockervolume = run_powershell(f"Get-BitLockerVolume -MountPoint {systemdrive}")
    protection_status = bitlockervolume.get("ProtectionStatus", 0)  # ProtectionStatus: Off = 0; On = 1
    volumestatus = volumestatus_dict[bitlockervolume.get("VolumeStatus", 0)]

    if tpm_state == "OK":
        protection_status = bitlockervolume.get("ProtectionStatus", 0)  # ProtectionStatus: Off = 0; On = 1
        volumestatus = volumestatus_dict[bitlockervolume.get("VolumeStatus", 0)]

        # Checking if drive is encrypted
        if protection_status == 1:
            print(f"OK: {systemdrive} OperatingSystem drive is encrypted and Protection status is ON")

            # Checking EncryptionMethod
            drive_encryption_method = bitlockervolume.get("EncryptionMethod", 0)
            if drive_encryption_method != target_encryption_method:
                print(
                    f"INFO: BitLocker encryption method is: {encryptionmethod_dict[drive_encryption_method]} and may should be: {encryptionmethod_dict[target_encryption_method]}"
                )
            else:
                print(f"INFO: BitLocker encryption method is: {encryptionmethod_dict[drive_encryption_method]}")

        # Checking if drive is encypted but protection status is OFF
        elif protection_status == 0 and volumestatus == "FullyEncrypted":
            print(f"WARNING: {systemdrive} OperatingSystem drive is encrypted but protection status is OFF")

            # Checking EncryptionMethod
            drive_encryption_method = bitlockervolume.get("EncryptionMethod", 0)
            if drive_encryption_method != target_encryption_method:
                print(
                    f"INFO: BitLocker encryption method is: {encryptionmethod_dict[drive_encryption_method]} and may should be: {encryptionmethod_dict[target_encryption_method]}"
                )
            else:
                print(f"INFO: BitLocker encryption method is: {encryptionmethod_dict[drive_encryption_method]}")

            # Resume Protection
            print(f"Resume Protection: {systemdrive} drive with BitLocker encryption method: {encryptionmethod_dict[drive_encryption_method]}")
            resume_bitlocker_pwsh = f"Resume-Bitlocker -MountPoint {systemdrive}"
            print(resume_bitlocker_pwsh)
            bitlockervolume = run_powershell(resume_bitlocker_pwsh)

            # Checking BitLocker state
            if bitlockervolume is not None:
                protection_status = bitlockervolume.get("ProtectionStatus", 0)  # ProtectionStatus: Off = 0; On = 1
                volumestatus = volumestatus_dict[bitlockervolume.get("VolumeStatus", 0)]
                if protection_status == 1:
                    print(f"BitLocker state of {systemdrive} drive is: {volumestatus} and Protection Status is ON")
                else:
                    print(f"BitLocker state of {systemdrive} drive is: {volumestatus} and Protection Status is OFF")
            else:
                if volumestatus != "FullyEncrypted" and force:
                    keyprotectorid = run_powershell(f"(Get-BitLockerVolume {systemdrive}).KeyProtector[0].KeyProtectorId")
                    remove_bitlockerkeyprotector_pwsh = f'Remove-BitlockerKeyProtector -MountPoint {systemdrive} -KeyProtectorId "{keyprotectorid}"'
                    print(remove_bitlockerkeyprotector_pwsh)
                    run_powershell(remove_bitlockerkeyprotector_pwsh)
                    error(f"BitlockerKeyProtector have been removed on {systemdrive} please reinstall this package.")
                else:
                    error(
                        "ERROR: The above PowerShell command appears to be unsuccessful.\nYou can force install this package to remove BitlockerKeyProtector."
                    )

        else:

            # Encrypting drive
            print(f"Encrypting: {systemdrive} drive with BitLocker encryption method: {encryptionmethod_dict[target_encryption_method]}")
            enable_bitlocker_pwsh = f"Enable-Bitlocker -MountPoint {systemdrive} -EncryptionMethod {encryptionmethod_dict[target_encryption_method]} -SkipHardwareTest -TpmProtector"
            print(enable_bitlocker_pwsh)

            if volumestatus != "EncryptionInProgress":
                bitlockervolume = run_powershell(enable_bitlocker_pwsh)
            else:
                print("Skipping the above PowerShell command because EncryptionInProgress.")

            # Checking BitLocker state
            if bitlockervolume is not None:
                volumestatus = volumestatus_dict[bitlockervolume.get("VolumeStatus", 0)]
                print(f"BitLocker state of {systemdrive} drive is: {volumestatus}")
            else:
                if volumestatus != "EncryptionInProgress" and force:
                    keyprotectorid = run_powershell(f"(Get-BitLockerVolume {systemdrive}).KeyProtector[0].KeyProtectorId")
                    remove_bitlockerkeyprotector_pwsh = f'Remove-BitlockerKeyProtector -MountPoint {systemdrive} -KeyProtectorId "{keyprotectorid}"'
                    print(remove_bitlockerkeyprotector_pwsh)
                    run_powershell(remove_bitlockerkeyprotector_pwsh)
                    error(f"BitlockerKeyProtector have been removed on {systemdrive} please reinstall this package.")
                else:
                    error(
                        "ERROR: The above PowerShell command appears to be unsuccessful.\nYou can force install this package to remove BitlockerKeyProtector."
                    )

    elif tpm_state == "WARNING":
        print("WARNING: This computer is not ready to silently enable BitLocker, you must fix his TPM chip state")

    elif tpm_state == "ERROR":
        print("ERROR: This computer cannot silently enable BitLocker, since no TPM chip have been found. You can only enable BitLocker manually")

    # Adding BitLocker Key Protector
    keyprotectortype_list = ",".join(
        run(f'powershell.exe -NoLogo -NoProfile -NonInteractive -ExecutionPolicy Unrestricted -OutputFormat text -Command "(Get-BitLockerVolume -MountPoint {systemdrive}).volumestatus"').splitlines()
        )
    if not "RecoveryPassword" in keyprotectortype_list:
        run(f'powershell.exe -NoLogo -NoProfile -NonInteractive -ExecutionPolicy Unrestricted -OutputFormat text -Command "Add-BitLockerKeyProtector -MountPoint {systemdrive} -RecoveryPasswordProtector"')


def audit():
    audit_status = "OK"
    is_package_installed = False

    for package in [p["package"] for p in WAPT.installed()]:
        if "audit-bitlocker" in package:
            is_package_installed = True
            break

    if not is_package_installed:
        print('WARNING: You have to install "tis-audit-bitlocker" to Backup-BitLockerKeyProtector in AD.')
        print('IMPORTANT: You have to create a customize package of "tis-audit-bitlocker" to Backup-BitLockerKeyProtector in WAPT.')
        audit_status = "WARNING"
    else:
        print("OK: BitLocker seems audited.")

    return audit_status


def check_tpm_state():
    audit_status = "OK"
    get_tpm = run_powershell("Get-Tpm")
    if get_tpm["TpmPresent"] == False:
        print("ERROR: No TPM chip found on this system")
        audit_status = "ERROR"
    else:
        print("OK: TPM chip found on this system")
        if get_tpm["TpmReady"] == False:
            print("WARNING: TPM chip not ready")
            audit_status = "WARNING"
        else:
            print("OK: TPM chip ready")
    WAPT.write_audit_data_if_changed("enable-bitlocker", "TpmPresent", get_tpm["TpmPresent"])
    WAPT.write_audit_data_if_changed("enable-bitlocker", "TpmReady", get_tpm["TpmReady"])
    if get_tpm["ManufacturerVersion"]:
        WAPT.write_audit_data_if_changed("enable-bitlocker", "ManufacturerVersion", get_tpm["ManufacturerVersion"].split("\x00")[0].strip())
    return audit_status