#!/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