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

try:
    import waptcrypto

    if "encrypted_data_str" in dir(waptcrypto):
        from waptcrypto import encrypted_data_str as rsa_encrypted_data_str
except:
    pass


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
decrypt_cert_list = []


def install():
    # Adding certificates allowed to decrypt in WAPT
    for cert in decrypt_cert_list:
        cert_path = makepath(WAPT.wapt_base_dir, "ssl", cert)
        if not isfile(cert_path):
            print("Copying: %s" % cert_path)
            filecopyto(cert, cert_path)


def audit():
    audit_status = "OK"
    # audit_status = check_tpm_state()

    certs_path = []
    for cert in decrypt_cert_list:
        certs_path.append(makepath(WAPT.wapt_base_dir, "ssl", cert))

    mountpoint_list = ensure_list(run_powershell("(Get-BitLockerVolume).MountPoint"))

    # Cleaning mountpoints (unpartitionned devices and removal devices)
    for mountpoint in mountpoint_list:
        is_ignored = False
        if "?" in mountpoint:
            print("INFO: An unknow volume has been detected and will be skipped (%s)" % mountpoint)
            is_ignored = True
        if run_powershell("Get-Volume -DriveLetter %s | Where-Object DriveType -EQ Removable" % mountpoint.replace(":", "")):
            is_ignored = True
        if is_ignored:
            mountpoint_list.remove(mountpoint)

    for mountpoint in mountpoint_list:
        bitlockervolume = run_powershell(f"Get-BitLockerVolume -MountPoint {mountpoint}")
        protection_status = bitlockervolume.get("ProtectionStatus", 0)  # ProtectionStatus: Off = 0; On = 1
        volumetype = "OperatingSystem" if bitlockervolume.get("VolumeType", 1) == 0 else "Data"  # VolumeType: OperatingSystem = 0; Data = 1
        volumestatus = volumestatus_dict[bitlockervolume.get("VolumeStatus", 0)]

        if protection_status == 1 and volumestatus =="FullyEncrypted":
            if volumetype == "OperatingSystem":
                print(f"OK: {mountpoint} {volumetype} drive is encrypted and protection status is ON with BitLocker")
                audit_status = set_audit_status(audit_status, "OK")
            else:
                print(f"OK: {mountpoint} {volumetype} drive is encrypted and protection status is ON with BitLocker")
                audit_status = set_audit_status(audit_status, "OK")  # Warning for external drives?

            try:
                keyprotector_list = ensure_list(run_powershell(f'(Get-BitLockerVolume -MountPoint "{mountpoint}").KeyProtector'))
            except:
                print(f"Unable to get KeyProtector informations for: {mountpoint} drive.")

            for keyprotector in keyprotector_list:
                keyprotectorid = keyprotector["KeyProtectorId"]
                if keyprotector["KeyProtectorType"] in [1,2,4]:
                    # WAPT.delete_audit_data(
                    #     "enable-bitlocker", f"RecoveryPassword_{keyprotectorid}"
                    # )  # not possible from package, please delete from WAPT Console
                    # print("INFO: Skipping Backup-BitLockerKeyProtector for KeyProtectorType: Tpm")
                    continue

                try:
                    # Backuping RecoveryPassword to the AD
                    if registry_readstring(HKEY_LOCAL_MACHINE, r'SOFTWARE\Microsoft\Windows\CurrentVersion\Group Policy\History', 'NetworkName', None):
                        print(f"Backuping: RecoveryPassword {keyprotectorid} to the AD")
                        run_powershell(f'Backup-BitLockerKeyProtector -MountPoint "{mountpoint}" -KeyProtectorId "{keyprotectorid}"')
                except Exception as e:
                    print(e)
                    print(f"WARNING: Failed to backup RecoveryPassword {keyprotectorid} to the AD")
                    audit_status = set_audit_status(audit_status, "WARNING")

                if decrypt_cert_list:
                    try:
                        # Backuping RecoveryPassword to the WAPT Console (encrypted)
                        print(f"Backuping: RecoveryPassword {keyprotectorid} to the WAPT Console")
                        WAPT.write_audit_data_if_changed(
                            "enable-bitlocker",
                            f"RecoveryPassword_{keyprotectorid}",
                            rsa_encrypted_data_str(keyprotector["RecoveryPassword"], certs_path),
                        )
                    except Exception as e:
                        print(e)
                        print(f"WARNING: Failed to backup RecoveryPassword {keyprotectorid} to the WAPT Console")
                        audit_status = set_audit_status(audit_status, "WARNING")
                else:
                    print(
                        f'IMPORTANT: You did not provided "decrypt_cert_list", so the RecoveryPassword {keyprotectorid} cannot be encrypted, so it will not be saved to the WAPT Console'
                    )

            keyprotectortype_list = ",".join(
                run('powershell.exe -NoLogo -NoProfile -NonInteractive -ExecutionPolicy Unrestricted -OutputFormat text -Command "&{(Get-BitLockerVolume -MountPoint "{mountpoint}").KeyProtector.KeyProtectorType}"').splitlines()
            )
            WAPT.write_audit_data_if_changed(
                "enable-bitlocker",
                "keyprotectortype_list_%s_drive" % mountpoint.split(":")[0],
                keyprotectortype_list,
            )
            continue

        elif volumestatus == "EncryptionInProgress":
            encryptpercent = bitlockervolume.get("EncryptionPercentage", 0)
            if volumetype == "OperatingSystem":
                print(f"WARNING: Encryption in progress : {encryptpercent}% on {mountpoint} {volumetype} drive")
                audit_status = set_audit_status(audit_status, "WARNING")
            else:
                print(f"WARNING: Encryption in progress : {encryptpercent}% on {mountpoint} {volumetype} drive")
                audit_status = set_audit_status(audit_status, "WARNING")  # Warning for external drives?
            continue

        elif volumestatus == "DecryptionInProgress":
            encryptpercent = bitlockervolume.get("EncryptionPercentage", 0)
            if volumetype == "OperatingSystem":
                print(f"WARNING: Decryption in progress : {encryptpercent}%  on {mountpoint} {volumetype} drive")
                audit_status = set_audit_status(audit_status, "WARNING")
            else:
                print(f"WARNING: Decryption in progress : {encryptpercent}%  on {mountpoint} {volumetype} drive")
                audit_status = set_audit_status(audit_status, "WARNING")  # Warning for external drives?
            continue

        elif protection_status == 0 and volumestatus =="FullyEncrypted":
            if volumetype == "OperatingSystem":
                print(f"ERROR: {mountpoint} {volumetype} drive is encrypted but protection status is OFF with BitLocker")
                audit_status = set_audit_status(audit_status, "ERROR")
            else:
                print(f"WARNING: {mountpoint} {volumetype} drive is encrypted but protection status is OFF with BitLocker")
                audit_status = set_audit_status(audit_status, "WARNING")  # Warning for external drives?
            continue

        else:
            print(f"ERROR: {mountpoint} {volumetype} drive is not encrypted with BitLocker")
            audit_status = set_audit_status(audit_status, "ERROR")

            # Checking EncryptionMethod
            drive_encryption_method = bitlockervolume.get("EncryptionMethod", 0)
            print(f"INFO: {mountpoint} {volumetype} drive BitLocker EncryptionMethod is: {encryptionmethod_dict[drive_encryption_method]}")


            WAPT.write_audit_data_if_changed(
                "enable-bitlocker",
                "encryptionmethod_%s_drive" % mountpoint.split(":")[0],
                encryptionmethod_dict[drive_encryption_method],
            )

    return audit_status


def set_audit_status(old_audit_status, new_audit_status):
    """Maintain higher criticality for audit status."""
    audit_level = {"OK": 0, "WARNING": 1, "ERROR": 2}

    old_status = old_audit_status.upper().strip()                 
    new_status = new_audit_status.upper().strip()

    if audit_level.get(new_status, -1) > audit_level.get(old_status, -1):
        return new_status
    else:
        return old_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