hidguardian.py

"""
This module provides objects to control the HIDGuardiand driver.
    This is achieved by modifying registry values that the driver looks for.

Classes can be used independantly or through the HidGuardian class
"""
import winreg

class AffectedDevices(object):
    """Manipulate a REG_Multi_SZ value to define HID devices

    Attributes:
        path (str): Path after HKLM to registry key containing AffectedDevices value.
        name (str): Name of value that contains our Affected Devices HID values.
    """
    path = r'SYSTEM\CurrentControlSet\Services\HidGuardian\Parameters'
    name = "AffectedDevices"

    def __init__(self):
        """When class object is created ensure that all registry path, keys, and values exist"""
        parent = r'SYSTEM\CurrentControlSet\Services\HidGuardian'

        winreg.CreateKeyEx(winreg.HKEY_LOCAL_MACHINE, parent, 0, winreg.KEY_WRITE)
        key = winreg.CreateKeyEx(winreg.HKEY_LOCAL_MACHINE, self.path, 0, winreg.KEY_ALL_ACCESS)

        try:
            winreg.QueryValueEx(key, self.name)
        except WindowsError:
            winreg.SetValueEx(key, self.name, 0, winreg.REG_MULTI_SZ, [])

        winreg.CloseKey(key)

    def list_devices(self):
        """List all HID devices registered in the AffectedDevices list
        Returns:
            list: list of strings
        """
        key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, self.path, 0, winreg.KEY_QUERY_VALUE)
        values = winreg.QueryValueEx(key, self.name)
        winreg.CloseKey(key)
        return values[0]

    def clear_devices(self):
        """Empties the list of AffectedDevices"""
        key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, self.path, 0, winreg.KEY_SET_VALUE)
        winreg.SetValueEx(key, self.name, 0, winreg.REG_MULTI_SZ, [])
        winreg.CloseKey(key)

    def add_device(self, hid):
        """Adds a device to the list of AffectedDevices
        Args:
            hid (str): Hardware ID for the desired device
        """
        values = self.list_devices()

        if hid not in values:
            values.append(hid)
            key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, self.path, 0, winreg.KEY_SET_VALUE)
            winreg.SetValueEx(key, self.name, 0, winreg.REG_MULTI_SZ, values)
            winreg.CloseKey(key)

    def remove_device(self, hid):
        """Removes a device from the list of AffectedDevices
        Args:
            hid (str): Hardware ID for the desired device
        """
        values = self.list_devices()

        if hid in values:
            while hid in values:
                values.remove(hid)
            key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, self.path, 0, winreg.KEY_SET_VALUE)
            winreg.SetValueEx(key, self.name, 0, winreg.REG_MULTI_SZ, values)
            winreg.CloseKey(key)

class WhiteList(object):
    """Manipulates a set of registry keys that reprisent Process Ids for the WhiteList
    Attributes:
        path (str): Path after HKLM to registry key containing the WhiteListed PID keys.
    """
    path = r'SYSTEM\CurrentControlSet\Services\HidGuardian\Parameters\WhiteList'

    def __init__(self):
        """When class object is created ensure that all registry path, and keys exist"""
        parent = r'SYSTEM\CurrentControlSet\Services\HidGuardian\Parameters'
        pparent = r'SYSTEM\CurrentControlSet\Services\HidGuardian'

        winreg.CreateKeyEx(winreg.HKEY_LOCAL_MACHINE, parent, 0, winreg.KEY_WRITE)
        winreg.CreateKeyEx(winreg.HKEY_LOCAL_MACHINE, pparent, 0, winreg.KEY_WRITE)
        winreg.CreateKeyEx(winreg.HKEY_LOCAL_MACHINE, self.path, 0, winreg.KEY_WRITE)

    def list_pids(self):
        """Enumarates the WhiteListed Process Ids
        Returns:
            list: a list of strings containing PIDs
        """
        key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, self.path, 0, winreg.KEY_ENUMERATE_SUB_KEYS)
        pids = []

        try:
            i = 0
            while True:
                pids.append(winreg.EnumKey(key, i))
                i += 1
        except WindowsError:
            pass

        winreg.CloseKey(key)
        return pids

    def clear_pids(self):
        """Deletes all WhiteListed PID keys to empty the WhiteList"""
        self.__delete_subkey(winreg.HKEY_LOCAL_MACHINE, self.path)
        winreg.CreateKeyEx(winreg.HKEY_LOCAL_MACHINE, self.path, 0, winreg.KEY_WRITE)

    def add_pid(self, pid):
        """Adds a PID to the WhiteList by creating a registry key
        Args:
            pid (str): Process ID for the program that needs to be WhiteListed
        """
        pids = self.list_pids()

        if pid not in pids:
            pkey = self.path + '\\' + pid
            winreg.CreateKeyEx(winreg.HKEY_LOCAL_MACHINE, pkey, 0, winreg.KEY_WRITE)

    def remove_pid(self, pid):
        """Removes a PID from the WhiteList by deleting a registry key
        Args:
            pid (str): Process ID for the program that needs to be Un-WhiteListed
        """
        pids = self.list_pids()

        if pid in pids:
            self.__delete_subkey(winreg.HKEY_LOCAL_MACHINE, self.path, pid)

    def __delete_subkey(self, key0, key1, key2=""):
        """Helper Function to delete a key and any subkey it may have
        Args:
            key0 (str): Registry Root such as HKEY_LOCAL_MACHINE
            key1 (str): Path to Parent Registry Key
            key2 (str): Path to registry key to be delete (along with its subkeys)
                If key2 is not supplied then key1 along with its subkeys will be deleted.
        """
        if key2 == "":
            currentkey = key1
        else:
            currentkey = key1 + "\\" + key2

        open_key = winreg.OpenKey(key0, currentkey, 0, winreg.KEY_ALL_ACCESS)
        infokey = winreg.QueryInfoKey(open_key)

        for dummy in range(0, infokey[0]):
            subkey = winreg.EnumKey(open_key, 0)
            try:
                winreg.DeleteKey(open_key, subkey)
            except WindowsError:
                self.__delete_subkey(key0, currentkey, subkey)

        winreg.DeleteKey(open_key, "")
        open_key.Close()
        return

class HidGuardian(object):
    """Object to enable, disable, and manipulate HIDGuardian
    Attributes:
        white_list (WhiteList): Associates a WhiteList object with itself
        affected_devices (AffectedDevices): Associates a Affected Devices object with itself
    """
    white_list = WhiteList()
    affected_devices = AffectedDevices()

    def __init__(self):
        """Check if driver is installed and/or update driver
        Args:
            inf (str): path to driver .inf file
            sys (str): path to driver .sys file

        """
        pass

    def start(self, pids, hids):
        """Start and Configure HIDGuardian
        Args:
            pids (list): list of strings containing desired PID values for the WhiteList
            hids (list): list of strings containing desired HID values for the AffectedDevices.
        """
        for pid in pids:
            self.white_list.add_pid(pid)

        for hid in hids:
            self.affected_devices.add_device(hid)

        # Start Driver / Reload device to activate filter

    def stop(self):
        """Stop and Clear Configuration for HIDGuardian"""
        self.affected_devices.clear_devices()
        self.white_list.clear_pids()

        # Reload / clear filter driver from Hid device