windowsclasses.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Place holder for all windows specific frequently used classes.
Currently here:
  RegistryHandler - Registry based operations
  DecideArchitecture - Accurate way to decide operating system architecture
"""

import os
import argparse
import errno
from _winreg import *  # noqa
try:
    from collections import *  # noqa
except ImporError:
    from OrderedDict import *  # noqa


def parse_arguments():
    """
    Argument parser, based on argparse module

    Keyword arguments:
    @return args -- arguments given by console
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-g", "--global",
        help="Whether we want to edit registry globally or just locally",
        action="store_true")
    parser.add_argument(
        "-d", "--different",
        help="Only handle the architecture specific registry area",
        action="store_true")
    parser.add_argument(
        "-r", "--registry",
        help="Which HKEY_* const registry type the program should use",
        type=str, choices=['HKLM', 'HKCR', 'HKCU', 'HKU', 'HKPD', 'HKCC'],
        default="HKLM")
    args = parser.parse_args()
    return args


def main():
    return RegistryHandler(parse_arguments())


class Struct:
    """
    Parameter bypassing struct
    """
    pass


class ClientRegistry:
    @staticmethod
    def directory():
        custom_param = Struct()
        custom_param.registry = "HKCU"
        handler = RegistryHandler(custom_param)
        directory = None
        try:
            directory = handler.get_key_value(
                "Software\CIRCLE Cloud Client", "running_directory")
        except LookupError:
            directory = os.path.dirname(os.path.abspath(__file__))
        return directory


class RegistryHandler:
    """
    Registry handling class, makes registry based queries and
    manipulations easier.
    This class can handle WOW64 based application differently and
    similarly (default)
    """

    def __init__(self, args=None):
        """Initialise RegistryHandler

        Keyword arguments:
        @param args  -- The arguments that decide how the we should
                        handle the registry
        args.registry    -- What base registry should we use
                            (default: HKLM)
        Raises AttributeError if not supported base registry type is
        given
        args.different -- Whether we handle registry redirection or not
        """
        if args is None:
            self.args = Struct()
            self.args.different = None
            self.args.registry = None
        else:
            if not hasattr(args, 'different'):
                args.different = None
            if not hasattr(args, 'registry'):
                args.registry = None
            self.args = args
        if self.args.different is None:
            self.args.different = False
        if self.args.registry is None or self.args.registry == "HKLM":
            self.args.registry = HKEY_LOCAL_MACHINE
        elif self.args.registry == "HKCR":
            self.args.registry = HKEY_CLASSES_ROOT
        elif self.args.registry == "HKCU":
            self.args.registry = HKEY_CURRENT_USER
        elif self.args.registry == "HKU":
            self.args.registry = HKEY_USERS
        elif self.args.registry == "HKPD":
            self.args.registry = HKEY_PERFORMANCE_DATA
        elif self.args.registry == "HKCC":
            self.args.registry = HKEY_CURRENT_CONFIG
        else:
            # print "Non supported registry type"
            raise AttributeError

    def connect_registry(self):
        """
        Getting a registry open

        Keyword arguments:
        @return connected_registy    -- Reference to the newly opened
                                        registry
        """
        return ConnectRegistry(None, self.args.registry)

    def create_registry_from_dict_chain(
            self, dict_chain, both=False, architect=KEY_WOW64_64KEY,
            needed_rights=KEY_ALL_ACCESS):
        """"
        Create registry key and value multilevel tree by chained
        dictionaries.
        Can raise AttributeError if the provided dict chain isn't
        correct

        Keyword arguments:
        @param key_value_chain -- The dict chain containing all the
                                  information
        @param both            -- Whether create the registry in WOW64
                                  node too
        @param architect       -- The current registry view (
                                  only change it if we want to fill the
                                  WOW6432 registry only [both = False],
                                  on x64 windows)
        @param needed_rights   -- SAM rights to access the key
        """
        if both and architect is KEY_WOW64_64KEY:
            self.create_registry_from_dict_chain(
                dict_chain, False, KEY_WOW64_32KEY, needed_rights)
        if not DecideArchitecture.Is64Windows():
            architect = 0
        connected_registy = self.connect_registry()
        if (isinstance(dict_chain, dict)
                or isinstance(dict_chain, OrderedDict)):
            for key, value in dict_chain.iteritems():
                if isinstance(value, dict) or isinstance(value, OrderedDict):
                    temp_dict = OrderedDict()
                    for my_key, my_value in value.iteritems():
                        temp_dict[key + "\\" + my_key] = my_value
                    self.create_registry_from_dict_chain(
                        temp_dict, False, architect, needed_rights)
                else:
                    if isinstance(value, list):
                        if len(value) % 2 != 0:
                            # print "Not enough member in the list"
                            raise AttributeError
                        else:
                            new_key = CreateKeyEx(
                                connected_registy, key, 0,
                                needed_rights | architect)
                            temp_dict = OrderedDict(
                                value[i:i + 2] for i in range(
                                    0, len(value), 2))
                            for my_key, my_value in temp_dict.iteritems():
                                if my_key == "default":
                                    my_key = None
                                SetValueEx(
                                    new_key, my_key, 0, REG_SZ, my_value)
                    else:
                        new_key = CreateKeyEx(
                            connected_registy, key, 0,
                            needed_rights | architect)
                        SetValueEx(new_key, None, 0, REG_SZ, value)
        else:
            print("The provided attribute wasn't a dictionary chain")
            raise AttributeError

    def get_key(self, key_name, needed_rights=KEY_ALL_ACCESS):
        """
        Getting a registry value by it's key's name
        Can raise KeyError if key is not found in the registry

        Keyword arguments:
        @param key_name      -- The specific key name of which value we
                                are interested in
        @param needed_rights -- SAM rights to access the key
        @return              -- [key, architect]
                                key          -- Reference to the opened
                                                key handler
                                architect    -- 0 for x86
                                                KEY_WOW64_32KEY or
                                                KEY_WOW64_64KEY
                                                depending where we
                                                found the key on x64
        """
        connected_registy = self.connect_registry()
        architect = 0
        if DecideArchitecture.Is64Windows():
            architect = KEY_WOW64_64KEY
        try:
            key = OpenKey(connected_registy, key_name, 0,
                          needed_rights | architect)
        except WindowsError:
            if DecideArchitecture.Is64Windows() and not self.args.different:
                try:
                    architect = KEY_WOW64_32KEY
                    key = OpenKey(connected_registy, key_name,
                                  0, needed_rights | architect)
                except WindowsError:
                    raise KeyError
            else:
                raise KeyError
        return [key, architect]

    def get_key_values(
            self, key_name, subkey_list, subroutine=False,
            depth="subkeys"):
        """
        Getting registry subkeys value by it's key's name and subkeys
        name
        Can raise LookupError exception if there are missing data
        Can raise AttributeError exception if depth attribute is wrong

        Keyword arguments:
        @param key_name      -- The specific key name of which subkey's
                                we are interested in
        @param subkey_list   -- List containing all the subkeys names
                                which values we are interested in
        @param subroutine    -- Whether suppress exception about not
                                having enough results or not
                                (default: False)
        @param depth         -- How depth the search should go for
                                [options: key, subkeys, all]
                                (default: subkeys)
        @return results{}    -- Dictionary with the subkey_name - value
                                combinations as keys and values
        """
        if depth == "key":
            int_depth = 0
        elif depth == "subkeys":
            int_depth = 1
        elif depth == "all":
            int_depth = 2
        else:
            raise AttributeError
        try:
            key_and_architect = self.get_key(key_name)
            key = key_and_architect[0]
            architect = key_and_architect[1]
        except KeyError:
            # print "%s doesn't exist in the registry" % key_name
            raise LookupError
        # print "%s found in the registry" % key_name
        results = {}
        if int_depth >= 1:
            for i in xrange(0, QueryInfoKey(key)[0] - 1):
                skey_name = EnumKey(key, i)
                skey = OpenKey(key, skey_name, 0, KEY_ALL_ACCESS | architect)
                if int_depth == 2 and QueryInfoKey(skey)[0] > 0:
                    for key, value in self.get_key_values(
                            skey_name, subkey_list, True, depth).iteritems():
                        results[key] = value
                for subkey_name in subkey_list:
                    try:
                        results[subkey_name] = QueryValueEx(
                            skey, subkey_name)[0]
                    except OSError as e:
                        if e.errno == errno.ENOENT:
                            # print ("%s doesn't exist in this" % subkey_name
                            #        " subkey")
                            pass
                skey.Close()
        if not results or len(results) != len(subkey_list):
            for subkey_name in subkey_list:
                try:
                    results[subkey_name] = QueryValueEx(key, subkey_name)[0]
                except OSError as e:
                    pass
        key.Close()
        if len(results) != len(subkey_list):
            # print "We are missing important variables"
            raise LookupError
        return results

    def get_key_value(
            self, key_name, subkey_name, subroutine=None, depth=None):
        """
        This is a wrapper for the get_key_values to be easier to use
        for single subkeys.
        Getting registry subkey value by it's key's name and subkey
        name

        Keyword arguments:
        @param key_name      -- The specific key name of which subkey's
                                we are interested in
        @param subkey_name   -- The specific subkey name which value we
                                are interested in
        @param subroutine    -- can be found at: get_key_values
        @param depth         -- can be found at: get_key_values
        @return value        -- Value of the specific subkey
        """
        try:
            if subroutine is None:
                return self.get_key_values(
                    key_name, [subkey_name])[subkey_name]
            elif depth is None:
                return self.get_key_values(
                    key_name, [subkey_name], subroutine)[subkey_name]
            else:
                return self.get_key_values(
                    key_name, [subkey_name], subroutine, depth)[subkey_name]
        except:
            raise


class DecideArchitecture:
    """
    Helper class to get the true ProgramFiles directory.
    This class doesn't depend on Phyton or Windows architecture.
    """

    @staticmethod
    def Is64Windows():
        return 'PROGRAMFILES(X86)' in os.environ

    @staticmethod
    def GetProgramFiles32():
        if DecideArchitecture.Is64Windows():
            return os.environ['PROGRAMFILES(X86)']
        else:
            return os.environ['PROGRAMFILES']

    @staticmethod
    def GetProgramFiles64():
        if DecideArchitecture.Is64Windows():
            return os.environ['PROGRAMW6432']
        else:
            return None