diff --git a/pupy/modules/reg.py b/pupy/modules/reg.py new file mode 100644 index 00000000..b4c18a8b --- /dev/null +++ b/pupy/modules/reg.py @@ -0,0 +1,225 @@ +# -*- coding: utf-8 -*- + +__class_name__ = 'reg' + +from pupylib.PupyModule import config, PupyModule, PupyArgumentParser +from pupylib.PupyOutput import Color, List, Table, MultiPart, TruncateToTerm + +TYPES = ( + 'NONE', 'SZ', 'EXPAND_SZ', 'BINARY', 'LE32', 'BE32', + 'LINK', 'MULTI_SZ', 'RESOURCE', 'RESOURCE_DESCRIPTOR', + 'RESOURCE_REQUIREMENTS_LIST' +) + +TYPE_COLORS = { + 'NONE': 'darkgrey', + 'SZ': 'white', + 'EXPAND_SZ': 'cyan', + 'BINARY': 'grey', + 'LE32': 'lightgreen', + 'BE32': 'blue', + 'LINK': 'lightyellow', + 'MULTI_SZ': 'lightred', + 'RESOURCE': 'red', + 'RESOURCE_DESCRIPTOR': 'red', + 'RESOURCE_REQUIREMENTS_LIST': 'red' +} + +def as_unicode(x): + if type(x) is str: + return x.decode('utf-8') + elif type(x) is unicode: + return x + else: + return unicode(x) + +@config(cat='admin', compatibilities=['windows']) +class reg(PupyModule): + '''Search/list/get/set/delete registry keys/values ''' + + dependencies = { + 'windows': ['reg'] + } + + @classmethod + def init_argparse(cls): + cls.arg_parser = PupyArgumentParser(prog='reg', description=cls.__doc__) + commands = cls.arg_parser.add_subparsers(dest='command') + + ls = commands.add_parser('ls') + ls.add_argument('key', default='HKCU', help='List key (HKCU by default)') + ls.add_argument('-w', '--wide', action='store_true', default=False, help='Show all the things') + ls.set_defaults(func=cls.ls) + + get = commands.add_parser('get') + get.add_argument('key', help='Key path (no default)') + get.add_argument( + 'name', nargs='?', default='', + help='Value name (default is Default value)') + get.set_defaults(func=cls.get) + + vset = commands.add_parser('set') + vset.add_argument( + '-c', '--create', action='store_true', default=False, + help='Create key') + vset.add_argument( + '-i', '--integer', action='store_true', default=False, + help='Set as DWORD/LE32, default SZ (string)') + vset.add_argument('key', help='Key path') + vset.add_argument('name', help='Value name to set') + vset.add_argument('value', help='Value to set') + vset.set_defaults(func=cls.set) + + rm = commands.add_parser('rm') + rm.add_argument('key', help='Key path') + rm.add_argument('name', default='', nargs='?', help='Value name or subkey to delete') + rm.set_defaults(func=cls.rm) + + search = commands.add_parser('search') + search.add_argument( + '-r', '--roots', default=('HKU', 'HKLM', 'HKCC'), nargs='+', help='Roots where to search') + search.add_argument( + '-K', '--exclude-key-name', action='store_false', default=True, + help='Do not search term in key names') + search.add_argument( + '-N', '--exclude-value-name', action='store_false', default=True, + help='Do not search term in value names') + search.add_argument( + '-V', '--exclude-value', action='store_false', default=True, + help='Do not search term in values') + search.add_argument( + '-R', '--regex', action='store_true', default=False, + help='Search term is regex') + search.add_argument( + '-E', '--equals', action='store_true', default=False, + help='Show only full matches') + search.add_argument( + '-i', '--ignorecase', action='store_true', default=False, + help='Ignore case') + search.add_argument( + '-w', '--wide', action='store_true', default=False, + help='Show all the things') + search.add_argument( + '-1', '--first', action='store_true', default=False, + help='Return after first match') + search.add_argument('term', help='Term to search') + search.set_defaults(func=cls.search) + + def run(self, args): + try: + args.func(self, args) + + except Exception, e: + if e.args[0] == 5: + self.error('Access denied') + elif e.args[0] == 'ascii': + self.error('Encoding error: {}'.format(e)) + else: + self.error('Error: {}'.format(e.args[0])) + + def _format_multi(self, results, wide=False, remove=None): + keys = [] + values = [] + + legend = ['NAME', 'TYPE', 'VALUE'] + if not remove: + legend.insert(0, 'KEY') + + for record in results: + is_key, key, rest = record[0], record[1], record[2:] + + if remove and key.startswith(remove): + key = key[len(remove)+1:] + + if is_key: + keys.append(key) + continue + + name, value, ktype = rest + + ktype = TYPES[ktype] + color = TYPE_COLORS[ktype] + + if not wide and type(value) in (str,unicode): + value = value.strip() + + values.append({ + 'KEY': Color(key, color), + 'NAME': Color(name, color), + 'VALUE': Color(value, color), + 'TYPE': Color(ktype, color) + }) + + results = [] + + if keys: + results.append(List(keys, caption='{ Keys }')) + + if values: + results.append(Table(values, legend, caption='Values')) + + if not keys and not values: + self.log('Empty') + else: + results = MultiPart(results) + if not wide: + results = TruncateToTerm(results) + self.log(results) + + def ls(self, args): + ls = self.client.remote('reg', 'enum') + result = ls(as_unicode(args.key)) + + if result is None: + self.error('No such key') + return + + self._format_multi(result, wide=args.wide, remove=args.key) + + def get(self, args): + get = self.client.remote('reg', 'get') + + value = get(as_unicode(args.key), as_unicode(args.name)) + if value is None: + self.error('No such key') + else: + self.log(value) + + def set(self, args): + kset = self.client.remote('reg', 'set') + value = args.value + + if args.integer: + value = int(value) + else: + value = as_unicode(value) + + try: + if kset(as_unicode(args.key), as_unicode(args.name), value, args.create): + self.success('OK') + else: + self.error('No such key') + except Exception, e: + print e + raise + + def rm(self, args): + rm = self.client.remote('reg', 'rm') + + if rm(as_unicode(args.key), as_unicode(args.name)): + self.success('OK') + else: + self.error('No such key') + + def search(self, args): + search = self.client.remote('reg', 'search') + results = search( + as_unicode(args.term), + tuple([as_unicode(x) for x in args.roots]), + args.exclude_key_name, args.exclude_value_name, + args.exclude_value, + args.regex, args.ignorecase, + args.first, args.equals + ) + + self._format_multi(results, wide=args.wide) diff --git a/pupy/packages/windows/all/reg.py b/pupy/packages/windows/all/reg.py new file mode 100644 index 00000000..cd8608c2 --- /dev/null +++ b/pupy/packages/windows/all/reg.py @@ -0,0 +1,426 @@ +# -*- coding: utf-8 -*- + +__all__ = [ + 'Key', 'search', 'enum', + 'set', 'get', 'rm' +] + +import _winreg +import re +import sys + +WELL_KNOWN_KEYS = { + 'HKEY_LOCAL_MACHINE': _winreg.HKEY_LOCAL_MACHINE, + 'HKLM': _winreg.HKEY_LOCAL_MACHINE, + 'HKEY_CURRENT_USER': _winreg.HKEY_CURRENT_USER, + 'HKCU': _winreg.HKEY_CURRENT_USER, + 'HKEY_CURRENT_CONFIG': _winreg.HKEY_CURRENT_CONFIG, + 'HKCC': _winreg.HKEY_CURRENT_CONFIG, + 'HKEY_CLASSES_ROOT': _winreg.HKEY_CLASSES_ROOT, + 'HKCR': _winreg.HKEY_CLASSES_ROOT, + 'HKEY_USERS': _winreg.HKEY_USERS, + 'HKU': _winreg.HKEY_USERS, + 'HKEY_PERFORMANCE_DATA': _winreg.HKEY_PERFORMANCE_DATA, + 'HKPD': _winreg.HKEY_PERFORMANCE_DATA, +} + +WELL_KNOWN_TYPES = { + int: _winreg.REG_DWORD, + str: _winreg.REG_SZ, + unicode: _winreg.REG_SZ, +} + +WELL_KNOWN_TYPES_NAMES = { + _winreg.REG_DWORD: 'DWORD', + _winreg.REG_BINARY: 'BINARY', + _winreg.REG_DWORD_LITTLE_ENDIAN: 'LE32', + _winreg.REG_DWORD_BIG_ENDIAN: 'BE32', + _winreg.REG_EXPAND_SZ: 'EXPAND_SZ', + _winreg.REG_LINK: 'LINK', + _winreg.REG_MULTI_SZ: 'MULTI_SZ', + _winreg.REG_NONE: 'NONE', + _winreg.REG_RESOURCE_LIST: 'RESOURCE', + _winreg.REG_FULL_RESOURCE_DESCRIPTOR: 'RESOURCE_DESCRIPTOR', + _winreg.REG_RESOURCE_REQUIREMENTS_LIST: 'RESOURCE_REQUIREMENTS_LIST', + _winreg.REG_SZ: 'SZ' +} + +class KeyIter(object): + __slots__ = ('handle', 'orig_name', 'key', 'sub', 'idx', 'is_value') + + def __init__(self, orig_name, key, sub, handle): + self.orig_name = orig_name + self.key = key + self.sub = sub + self.handle = handle + self.idx = 0 + self.is_value = False + + def next(self): + try: + result = None + if self.is_value: + name, value, ktype = _winreg.EnumValue(self.handle, self.idx) + result = Value(self.orig_name, name, value, ktype) + else: + value = _winreg.EnumKey(self.handle, self.idx) + value = value.decode(sys.getfilesystemencoding()) + result = Key(self.orig_name + '\\' + value) + + self.idx += 1 + return result + + except WindowsError, e: + if e.winerror != 259: + raise + + if self.is_value: + raise StopIteration() + else: + self.idx = 0 + self.is_value = True + return self.next() + +class Value(object): + __slots__ = ('parent', 'name', 'value', 'type') + + def __init__(self, parent, name, value, ktype): + if type(parent) == str: + parent = parent.decode(sys.getfilesystemencoding()) + + if type(name) == str: + name = name.decode(sys.getfilesystemencoding()) + + if type(value) == str and ktype in (_winreg.REG_SZ, _winreg.REG_MULTI_SZ): + value = value.decode(sys.getfilesystemencoding()) + + self.parent = parent + self.name = name + self.value = value + self.type = ktype + + def __str__(self): + if self.type == _winreg.REG_EXPAND_SZ: + value = self.value + if type(value) is not unicode: + try: + value = value.decode(sys.getfilesystemencoding()) + except UnicodeDecodeError: + value = value.decode('latin1') + + try: + return _winreg.ExpandEnvironmentStrings(value) + except TypeError: + pass + + return value + + elif self.type == _winreg.REG_SZ: + return self.value + + return str(self.value) + + def __int__(self): + return int(self.value) + + def __repr__(self): + return '{}={} ({})'.format( + self.name, self.value, WELL_KNOWN_TYPES_NAMES[self.type]) + +class Key(object): + __slots__ = ('arg', 'key', 'sub', 'access', 'create') + + def __init__(self, key, access=_winreg.KEY_READ, create=False): + sub_key = None + top_key = None + + if type(key) == str: + key = key.decode(sys.getfilesystemencoding()) + + for wkk, wrk in WELL_KNOWN_KEYS.iteritems(): + if key == wkk: + top_key = wrk + sub_key = '' + elif key.startswith((wkk+'\\', wkk+'/')): + top_key = wrk + sub_key = key[len(wkk)+1:] + break + + if not top_key: + raise KeyError(key) + + sub_key = sub_key.strip('\\') + + self.key = top_key + self.sub = sub_key + self.arg = key + self.access = access + self.create = create + + @property + def name(self): + return self.arg.split('\\')[-1] + + @property + def parent(self): + return '\\'.join(self.arg.split('\\')[:-1]) + + def _open_key(self, access=None): + try: + sub = self.sub + + if type(sub) == unicode: + sub = sub.encode(sys.getfilesystemencoding()) + + if self.create: + return _winreg.CreateKey( + self.key, sub) + else: + return _winreg.OpenKey( + self.key, sub, 0, + access if access is not None else self.access) + except WindowsError, e: + if e.winerror != 2: + raise + + raise KeyError(self.sub) + + def _query_value(self, handle, attr): + try: + if type(attr) == unicode: + attr = attr.encode(sys.getfilesystemencoding()) + value, ktype = _winreg.QueryValueEx(handle, attr) + except WindowsError, e: + if e.winerror != 2: + raise + + raise KeyError(attr) + + return Value(self.arg, attr, value, ktype) + + def __iter__(self): + handle = self._open_key( + _winreg.KEY_QUERY_VALUE | _winreg.KEY_ENUMERATE_SUB_KEYS) + + iterator = KeyIter(self.arg, self.key, self.sub, handle) + try: + while True: + try: + yield next(iterator) + except WindowsError, e: + ## Ignore everything + print "PISKA", e.args + pass + + except StopIteration: + return + + finally: + _winreg.CloseKey(handle) + + def __str__(self): + return self.arg.encode(sys.getfilesystemencoding()) + + def __unicode__(self): + return self.arg + + def __repr__(self): + return self.arg + + def __delitem__(self, attr): + handle = self._open_key(_winreg.KEY_SET_VALUE) + + if type(attr) == unicode: + attr = attr.encode(sys.getfilesystemencoding()) + + try: + try: + return _winreg.DeleteValue(handle, attr) + except WindowsError, e: + if e.winerror != 2: + raise + + try: + return _winreg.DeleteKey(handle, attr) + except WindowsError, e: + if e.winerror != 2: + raise + + raise KeyError(attr) + + finally: + _winreg.CloseKey(handle) + + def __getitem__(self, attr): + handle = self._open_key() + try: + return self._query_value(handle, attr) + finally: + _winreg.CloseKey(handle) + + def __setitem__(self, attr, value): + vtype = type(value) + + if vtype not in WELL_KNOWN_TYPES and vtype is not Value: + raise TypeError('setattr only supported for str/int') + + ktype = value.type if vtype is Value else WELL_KNOWN_TYPES[vtype] + if vtype is Value: + value = value.value + + handle = self._open_key(_winreg.KEY_SET_VALUE) + + if type(attr) == unicode: + attr = attr.encode(sys.getfilesystemencoding()) + + if type(value) == unicode: + value = value.encode(sys.getfilesystemencoding()) + + try: + return _winreg.SetValueEx( + handle, attr, 0, ktype, value) + finally: + _winreg.CloseKey(handle) + +def search(term, roots=('HKU', 'HKLM', 'HKCC'), key=True, name=True, value=True, regex=False, ignorecase=False, first=False, equals=False): + compare = None + + def as_str(x): + vtype = type(x) + if vtype in (str, unicode): + return x + else: + return str(x) + + def contains(x, y): + try: + if type(x) == str and type(y) == unicode: + return x in y.encode(sys.getfilesystemencoding()) + elif type(x) == unicode and type(y) == str: + return x.encode(sys.getfilesystemencoding()) in y.encode + else: + return x in y + + except Exception: + return False + + if regex: + term = re.compile( + term, + re.MULTILINE | (re.IGNORECASE if ignorecase else 0)) + + if equals: + compare = lambda x: term.match(as_str(x)) + else: + compare = lambda x: term.search(as_str(x)) + elif ignorecase: + if equals: + compare = lambda x: term.lower() == as_str(x).lower() + else: + compare = lambda x: term.lower() in as_str(x).lower() + else: + if equals: + compare = lambda x: term == x + else: + compare = lambda x: contains(term, x) + + if type(roots) in (str, unicode): + roots = [roots] + + def _walk(root): + results = [] + + for kv in root: + tkv = type(kv) + + if tkv == Key: + if key and compare(kv.name): + results.append(kv) + if first and results: + break + + try: + results.extend(_walk(kv)) + + if first and results: + break + + except WindowsError: + pass + + elif tkv == Value: + if name and compare(kv.name): + results.append(kv) + + elif value and (equals or type(kv.value) in (str, unicode)) and compare(kv.value): + results.append(kv) + + if first and results: + break + + else: + raise TypeError('Unknown type {} in search'.format(tkv.__name__)) + + return results + + typleized = [] + for root in roots: + for result in _walk(Key(root)): + if type(result) == Key: + typleized.append((True, str(result))) + else: + typleized.append(( + False, result.parent, result.name, + result.value, result.type)) + + return typleized + +def enum(path): + try: + tupleized = [] + + for item in Key(path): + if type(item) == Key: + tupleized.append((True, unicode(item))) + else: + tupleized.append(( + False, item.parent, item.name, + item.value, item.type)) + + return tuple(tupleized) + + except KeyError: + return None + +def set(path, name, value, create): + try: + k = Key(path, create=create) + try: + old_value = k[name] + if old_value.type in (_winreg.REG_SZ, _winreg.REG_BINARY): + if type(value) not in (str, unicode): + value = str(value) + elif old_value.type in (_winreg.REG_DWORD, _winreg.REG_DWORD_LITTLE_ENDIAN): + if type(value) != int: + value = int(value) + except KeyError: + pass + + k[name] = value + return True + + except KeyError: + return False + +def get(path, name): + try: + return Key(path)[name].value + except KeyError: + return None + +def rm(path, name): + try: + del Key(path)[name] + return True + except KeyError: + return False