Add render method to registry

This commit is contained in:
Mathieu Comandon 2016-10-25 14:38:29 -07:00
parent ca9de01538
commit fa67361a68
2 changed files with 43 additions and 6 deletions

View file

@ -50,9 +50,13 @@ class WindowsFileTime:
class WineRegistry(object):
version_header = "WINE REGISTRY Version "
relative_to_header = ";; All keys relative to "
def __init__(self, reg_filename=None):
self.arch = 'win32'
self.version = 2
self.relative_to = "\\\\User\\\\S-1-5-21-0-0-0-1000"
self.keys = OrderedDict()
if reg_filename:
self.prefix_path = os.path.dirname(reg_filename)
@ -68,21 +72,38 @@ class WineRegistry(object):
registry_lines = self.get_raw_registry(reg_filename)
current_key = None
for line in registry_lines:
# Remove trailing newlines
line = line.rstrip('\n')
if line.startswith('WINE REGISTRY Version'):
self.version = int(line[len('WINE REGISTRY Version '):])
line = line.rstrip('\n') # Remove trailing newlines
if line.startswith(self.version_header):
self.version = int(line[len(self.version_header):])
continue
if line.startswith(self.relative_to_header):
self.relative_to = line[len(self.relative_to_header):]
continue
if line.startswith('#arch'):
self.arch = line.split('=')[1]
continue
if line.startswith('['):
current_key = WineRegistryKey(key_def=line)
self.keys[current_key.name] = current_key
continue
if current_key:
current_key.parse(line)
def render(self):
content = ""
content += "{}{}\n".format(self.version_header, self.version)
content += "{}{}\n\n".format(self.relative_to_header, self.relative_to)
content += "#arch={}\n".format(self.arch)
for key in self.keys:
content += "\n"
content += self.keys[key].render()
return content
def query(self, keypath, value=None):
key = self.keys.get(keypath)
if key:

View file

@ -7,8 +7,8 @@ FIXTURES_PATH = os.path.join(os.path.dirname(__file__), 'fixtures')
class TestWineRegistry(TestCase):
def setUp(self):
registry_path = os.path.join(FIXTURES_PATH, 'user.reg')
self.registry = WineRegistry(registry_path)
self.registry_path = os.path.join(FIXTURES_PATH, 'user.reg')
self.registry = WineRegistry(self.registry_path)
def test_can_load_registry(self):
self.assertTrue(len(self.registry.keys) > 10)
@ -48,3 +48,19 @@ class TestWineRegistry(TestCase):
)
key = self.registry.keys.get('Software/Wine/Fonts')
self.assertEqual(key.render(), expected)
def test_render_user_reg(self):
content = self.registry.render()
with open(self.registry_path, 'r') as registry_file:
original_content = registry_file.read()
self.assertEqual(content, original_content)
def test_can_render_system_reg(self):
registry_path = os.path.join(FIXTURES_PATH, 'system.reg')
with open(registry_path, 'r') as registry_file:
original_content = registry_file.read()
system_reg = WineRegistry(registry_path)
content = system_reg.render()
with open(registry_path + '.new', 'w') as new_file:
new_file.write(content)
self.assertEqual(content, original_content)