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

# Usable WAPT package functions: install(), uninstall(), session_setup(), audit(), update_package()
# Declaring global variables - Warnings: 1) WAPT context is only available in package functions; 2) Global variables are not persistent between calls


def install():
    pass


def audit():
    # Declaring local variables
    return_warning = False
    return_error = False

    res_tpm = run_powershell("Get-Tpm")
    if res_tpm['TpmPresent'] == False:
        print("ERROR: No TPM chip found on this system")
        return "ERROR"
    else:
        print("OK: TPM chip found on this system")
        if res_tpm['TpmReady'] == True:
            print('OK: TPM chip ready')
        else:
            print("WARNING: TPM chip not ready")
            return_warning = True

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

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

    for mountpoint in mountpoint_list:
        try:
            recovery_key = run_powershell("(Get-BitLockerVolume -MountPoint {}).KeyProtector.KeyProtectorId".format(mountpoint))
            if recovery_key:
                print("OK: {} recovery keys found".format(mountpoint))
                for mountpoint in mountpoint_list:
                    try:
                        for key_protectorid in run_powershell('(Get-BitLockerVolume -MountPoint "{}").KeyProtector.KeyProtectorId'.format(mountpoint)):
                            print("Storing: {} recovery key to the AD".format(key_protectorid))
                            run_powershell('Backup-BitLockerKeyProtector -MountPoint "{}" -KeyProtectorId "{}"'.format(mountpoint, key_protectorid))
                    except Exception as e:
                        print("WARNING: Failed to store recovery keys to the AD for: {}".format(mountpoint))
                        return_warning = True
        except:
            print("WARNING: Failed to recover {} mountpoint key".format(mountpoint))
            return_warning = True

    for mountpoint in mountpoint_list:
        if run_powershell('Get-BitLockerVolume -MountPoint %s | Select-Object -ExpandProperty ProtectionStatus' % mountpoint) == 0:
            print("ERROR: %s Drive is NOT encrypted with BitLocker" % mountpoint)
            return_error = True
        else:
            print("OK: %s Drive is encrypted with BitLocker" % mountpoint)

    if return_error:
        return "ERROR"

    if return_warning:
        return "WARNING"

    return "OK"


def update_package():
    pass