# Wizard Kit: Functions - Cleanup

from functions.common import *

# STATIC VARIABLES
D7_HKCR_CLEANUP = {
    r'batfile\shell\{0001B4FD-9EA3-4D90-A79E-FD14BA3AB01D}': {'Recurse': True},
    r'cmdfile\shell\{0001B4FD-9EA3-4D90-A79E-FD14BA3AB01D}': {'Recurse': True},
    r'exefile\shell\ResourceHacker': {'Recurse': True},
    r'regfile\shell\{0001B4FD-9EA3-4D90-A79E-FD14BA3AB01D}': {'Recurse': True},
    }
D7_HKCU_CLEANUP = {
    r'Software\Malwarebytes': {'Recurse': False},
    }
D7_HKLM_CLEANUP = {
    r'Software\Emsisoft': {'Recurse': False},
    }
HKU =  winreg.HKEY_USERS
HKCR = winreg.HKEY_CLASSES_ROOT
HKCU = winreg.HKEY_CURRENT_USER
HKLM = winreg.HKEY_LOCAL_MACHINE

def cleanup_adwcleaner():
    """Move AdwCleaner folders into the ClientDir."""
    source_path = r'{SYSTEMDRIVE}\AdwCleaner'.format(**global_vars['Env'])
    source_quarantine = r'{}\Quarantine'.format(source_path)

    # Quarantine
    if os.path.exists(source_quarantine):
        os.makedirs(global_vars['QuarantineDir'], exist_ok=True)
        dest_name = r'{QuarantineDir}\AdwCleaner_{Date-Time}'.format(
            **global_vars)
        dest_name = non_clobber_rename(dest_name)
        shutil.move(source_quarantine, dest_name)

    # Delete source folder if empty
    try:
        os.rmdir(source_path)
    except OSError:
        pass

    # Main folder
    if os.path.exists(source_path):
        os.makedirs(global_vars['LogDir'], exist_ok=True)
        dest_name = r'{LogDir}\{Date}\AdwCleaner'.format(
            **global_vars)
        dest_name = non_clobber_rename(dest_name)
        shutil.move(source_path, dest_name)

def cleanup_cbs(dest_folder):
    """Safely cleanup a known CBS archive bug under Windows 7.

    If a CbsPersist file is larger than 2 Gb then the auto archive feature
    continually fails and will fill up the system drive with temp files.

    This function moves the temp files and CbsPersist file to a temp folder,
    compresses the CbsPersist files with 7-Zip, and then opens the temp folder
    for the user to manually save the backup files and delete the temp files.
    """
    backup_folder = r'{dest_folder}\CbsFix'.format(dest_folder=dest_folder)
    temp_folder = r'{backup_folder}\Temp'.format(backup_folder=backup_folder)
    os.makedirs(backup_folder, exist_ok=True)
    os.makedirs(temp_folder, exist_ok=True)

    # Move files into temp folder
    cbs_path = r'{SYSTEMROOT}\Logs\CBS'.format(**global_vars['Env'])
    for entry in os.scandir(cbs_path):
        # CbsPersist files
        if entry.name.lower().startswith('cbspersist'):
            dest_name = r'{}\{}'.format(temp_folder, entry.name)
            dest_name = non_clobber_rename(dest_name)
            shutil.move(entry.path, dest_name)
    temp_path = r'{SYSTEMROOT}\Temp'.format(**global_vars['Env'])
    for entry in os.scandir(temp_path):
        # cab_ files
        if entry.name.lower().startswith('cab_'):
            dest_name = r'{}\{}'.format(temp_folder, entry.name)
            dest_name = non_clobber_rename(dest_name)
            shutil.move(entry.path, dest_name)

    # Compress CbsPersist files with 7-Zip
    cmd = [
        global_vars['Tools']['SevenZip'],
        'a', '-t7z', '-mx=3', '-bso0', '-bse0',
        r'{}\CbsPersists.7z'.format(backup_folder),
        r'{}\CbsPersist*'.format(temp_folder)]
    run_program(cmd)

def cleanup_d7ii():
    """Sort d7II logs and remove temp items."""
    d7_path = r'{}\d7II'.format(global_vars['ClientDir'])
    d7_reports = r'{}_Reports'.format(d7_path)
    d7_temp = r'{}\Temp'.format(d7_path)

    # Logs & Reports
    if os.path.exists(d7_reports):
        for entry in os.scandir(d7_reports):
            r = re.match(r'(\d+)-(\d+)-(\d+)', entry.name)
            d7_date = '{}-{:02d}-{:02d}'.format(
                r.group(1), int(r.group(2)), int(r.group(3)))
            d7_mlogs = r'{}\Malware Logs'.format(entry.path)
            log_dest = r'{SYSTEMDRIVE}\{prefix}\Info\{date}'.format(
                prefix=KIT_NAME_SHORT,
                date=d7_date,
                **global_vars['Env'])

            # Remove empty folders
            for f in ('Malware Logs', 'Screen Shots'):
                try:
                    os.rmdir(r'{}\{}'.format(entry.path, f))
                except FileNotFoundError:
                    pass
                except OSError:
                    pass

            # Malware Logs
            if os.path.exists(d7_mlogs):
                for m_entry in os.scandir(d7_mlogs):
                    prefix = ''
                    if m_entry.name == 'MalwareScan_Report.txt':
                        prefix = 'd7II_'
                    dest_path = r'{log_dest}\{prefix}{name}'.format(
                        log_dest=log_dest,
                        prefix=prefix,
                        name=m_entry.name)
                    dest_path = non_clobber_rename(dest_path)
                    shutil.move(entry.path, dest_path)
                try:
                    os.rmdir(d7_mlogs)
                except OSError:
                    pass

            # Other items
            for o_entry in os.scandir(entry.path):
                dest_path = r'{log_dest}\d7II_{name}'.format(
                    log_dest=log_dest,
                    name=m_entry.name)
                dest_path = non_clobber_rename(dest_path)
                shutil.move(entry.path, dest_path)

            # Remove folder if empty
            try:
                os.rmdir(entry.path)
            except OSError:
                pass

    # Registry Items
    for key, settings in D7_HKCR_CLEANUP.items():
        delete_registry_key(HKCR, key, recurse=settings['Recurse'])
    for key, settings in D7_HKCU_CLEANUP.items():
        delete_registry_key(HKCU, key, recurse=settings['Recurse'])
    for key, settings in D7_HKLM_CLEANUP.items():
        delete_registry_key(HKLM, key, recurse=settings['Recurse'])

    # Temp items
    if os.path.exists(d7_path):
        if os.path.exists(d7_temp):
            shutil.rmtree(d7_temp)
        try:
            os.rmdir(d7_path)
        except OSError:
            pass

def cleanup_desktop():
    """Move known backup files and reports into the ClientDir."""
    dest_folder = r'{ProgBackupDir}\{Date}\Desktop'.format(**global_vars)
    os.makedirs(dest_folder, exist_ok=True)

    desktop_path = r'{USERPROFILE}\Desktop'.format(**global_vars['Env'])
    for entry in os.scandir(desktop_path):
        # JRT, RKill, Shortcut cleaner
        if re.search(r'^(JRT|RKill|sc-cleaner)', entry.name, re.IGNORECASE):
            dest_name = r'{}\{}'.format(dest_folder, entry.name)
            dest_name = non_clobber_rename(dest_name)
            shutil.move(entry.path, dest_name)

    # Remove dir if empty
    try:
        os.rmdir(dest_folder)
    except OSError:
        pass

def cleanup_emsisoft():
    """Remove EmsisoftCmd files from drive root."""
    source_path = r'{}\EmsisoftCmd'.format(global_vars['Env']['SYSTEMDRIVE'])
    source_quarantine = r'{}\Quarantine'.format(source_path)

    # Quarantine
    if os.path.exists(source_quarantine):
        os.makedirs(global_vars['QuarantineDir'], exist_ok=True)
        dest_name = r'{QuarantineDir}\Emsisoft_{Date-Time}'.format(
            **global_vars)
        dest_name = non_clobber_rename(dest_name)
        shutil.move(source_quarantine, dest_name)

    # Remove program
    if os.path.exists(source_path):
        shutil.rmtree(source_path)

def cleanup_regbackups():
    """Move d7ii regbackups into backup folder."""
    source_path = r'{}\Support\RegBackups'.format(
        global_vars['Env']['SYSTEMDRIVE'])

    # Bail early
    if not os.path.exists(source_path):
        return

    # Move to backup folder
    for entry in os.scandir(source_path):
        os.makedirs(global_vars['ProgBackupDir'], exist_ok=True)
        dest_path = r'{ProgBackupDir}\{Date}\Registry\{name}'.format(
            name=entry.name,
            **global_vars)
        dest_path = non_clobber_rename(dest_path)
        shutil.move(entry.path, dest_path)

    # Delete source folders if empty
    try:
        os.rmdir(source_path)
        os.rmdir(r'{}\Support'.format(global_vars['Env']['SYSTEMDRIVE']))
    except OSError:
        pass

def delete_registry_key(hive, key, recurse=False):
    """Delete a registry key and all it's subkeys."""
    access = winreg.KEY_ALL_ACCESS

    if recurse:
        # Delete all subkeys first
        with winreg.OpenKeyEx(hive, key, 0, access) as k:
            key_info = winreg.QueryInfoKey(k)
            for x in range(key_info[0]):
                subkey = r'{}\{}'.format(key, winreg.EnumKey(k, 0))
                delete_registry_key(hive, subkey)

    # Delete key
    winreg.DeleteKey(hive, key)

def delete_registry_value(hive, key, value):
    """Delete a registry value."""
    access = winreg.KEY_ALL_ACCESS
    with winreg.OpenKeyEx(hive, key, 0, access) as k:
        winreg.DeleteValue(k, value)

if __name__ == '__main__':
    print("This file is not meant to be called directly.")

# vim: sts=4 sw=4 ts=4