diff --git a/README.rst b/README.rst index 363f9f16..057e63f7 100644 --- a/README.rst +++ b/README.rst @@ -39,13 +39,6 @@ Simple Example # save database >>> kp.save() -Context Manager Example --------------- -.. code:: python - >>> with PyKeePass('db.kdbx', password='somePassw0rd') as kp: - >>> entry = kp.find_entries(title='facebook', first=True) - >>> entry.password - 's3cure_p455w0rd' Finding Entries ---------------------- @@ -173,13 +166,13 @@ For backwards compatibility, the following functions are also available: Adding Entries -------------- -**add_entry** (destination_group, title, username, password, url=None, notes=None, tags=None, expiry_time=None, icon=None, force_creation=False) +**add_entry** (destination_group, title, username, password, url=None, notes=None, tags=None, expiry_time=None, icon=None, customicon=None, force_creation=False) **delete_entry** (entry) **move_entry** (entry, destination_group) -where ``destination_group`` is a ``Group`` instance. ``entry`` is an ``Entry`` instance. ``title``, ``username``, ``password``, ``url``, ``notes``, ``tags``, ``icon`` are strings. ``expiry_time`` is a ``datetime`` instance. +where ``destination_group`` is a ``Group`` instance. ``entry`` is an ``Entry`` instance. ``title``, ``username``, ``password``, ``url``, ``notes``, ``tags``, ``icon``, ``customicon`` are strings. ``expiry_time`` is a ``datetime`` instance. If ``expiry_time`` is a naive datetime object (i.e. ``expiry_time.tzinfo`` is not set), the timezone is retrieved from ``dateutil.tz.gettz()``. @@ -189,9 +182,9 @@ If ``expiry_time`` is a naive datetime object (i.e. ``expiry_time.tzinfo`` is no >>> kp.add_entry(kp.root_group, 'testing', 'foo_user', 'passw0rd') Entry: "testing (foo_user)" - # add a new entry to the social group + # add a new entry to the social group with a custom icon >>> group = find_groups(name='social', first=True) - >>> entry = kp.add_entry(group, 'testing', 'foo_user', 'passw0rd') + >>> entry = kp.add_entry(group, 'testing', 'foo_user', 'passw0rd', customicon="2") Entry: "testing (foo_user)" # save the database diff --git a/pykeepass/baseelement.py b/pykeepass/baseelement.py index 2bda4032..12ae9137 100644 --- a/pykeepass/baseelement.py +++ b/pykeepass/baseelement.py @@ -1,7 +1,7 @@ from __future__ import unicode_literals from __future__ import absolute_import from lxml import etree -from lxml.etree import Element +from lxml.etree import Element, iterwalk from lxml.builder import E from datetime import datetime, timedelta import base64 @@ -13,15 +13,18 @@ class BaseElement(object): """Entry and Group inherit from this class""" - def __init__(self, element=None, version=None, icon=None, expires=False, - expiry_time=None): + def __init__(self, element=None, version=None, icon=None, customicon=None, + expires=False, meta=None, expiry_time=None): self._element = element + self._meta = meta self._element.append( E.UUID(base64.b64encode(uuid.uuid1().bytes).decode('utf-8')) ) if icon: self._element.append(E.IconID(icon)) + if customicon: + self._element.append(E.CustomIconUUID(self._get_customicon_uuid(customicon))) current_time_str = self._encode_time(datetime.utcnow()) if expiry_time: expiry_time_str = self._encode_time( @@ -51,7 +54,8 @@ def _set_subelement_text(self, tag, value): v = self._element.find(tag) if v is not None: self._element.remove(v) - self._element.append(getattr(E, tag)(value)) + if value is not None: + self._element.append(getattr(E, tag)(value)) def dump_xml(self, pretty_print=False): return etree.tostring(self._element, pretty_print=pretty_print) @@ -72,6 +76,14 @@ def icon(self): def icon(self, value): return self._set_subelement_text('IconID', value) + @property + def customicon(self): + return self._get_customicon_id(self._get_subelement_text('CustomIconUUID')) + + @customicon.setter + def customicon(self, value): + return self._set_subelement_text('CustomIconUUID', self._get_customicon_uuid(value)) + @property def _path(self): return self._element.getroottree().getpath(self._element) @@ -121,7 +133,6 @@ def _decode_time(self, text): tzinfos={'UTC':tz.gettz('UTC')} ) - def _get_times_property(self, prop): times = self._element.find('Times') if times is not None: @@ -136,6 +147,25 @@ def _set_times_property(self, prop, value): if prop is not None: prop.text = self._encode_time(value) + def _get_customicon_uuid(self, iconid): + iconid = int(iconid) + if self._meta is not None: + icons = self._meta.xpath('.//CustomIcons/Icon/UUID[1]') + return str(icons[iconid].text) if iconid >= 0 and iconid < len(icons) else None + + def _get_customicon_id(self, iconuuid): + if self._meta is not None: + walk = iterwalk(self._meta, tag="CustomIcons") + find = self._meta.xpath('.//Icon/UUID[text()="{}"]'.format(iconuuid)) + if len(find) >= 1: + icon = find[0].getparent() + for _, element in walk: + icons = element.getchildren() + iconid = icons.index(icon) + self.icon = "0" + return str(iconid) + break + @property def expires(self): times = self._element.find('Times') @@ -152,7 +182,6 @@ def expires(self, value): def expired(self): return self.expires and (datetime.utcnow() > self.expiry_time) - @property def expiry_time(self): return self._get_times_property('ExpiryTime') diff --git a/pykeepass/entry.py b/pykeepass/entry.py index 2dcd533a..6d3470b2 100644 --- a/pykeepass/entry.py +++ b/pykeepass/entry.py @@ -2,7 +2,7 @@ from __future__ import absolute_import from pykeepass.baseelement import BaseElement from copy import deepcopy -from lxml.etree import Element, _Element +from lxml.etree import Element, _Element, iterwalk from lxml.objectify import ObjectifiedElement from lxml.builder import E import logging @@ -17,6 +17,7 @@ 'URL', 'Tags', 'IconID', + 'CustomIconUUID', 'Times', 'History' ] @@ -26,12 +27,13 @@ class Entry(BaseElement): def __init__(self, title=None, username=None, password=None, url=None, notes=None, tags=None, expires=False, expiry_time=None, - icon=None, element=None, version=None): + icon=None, customicon=None, element=None, version=None, meta=None): assert type(version) is tuple, 'The provided version is not a tuple, but a {}'.format( type(version) ) self._version = version + self._meta = meta if element is None: super(Entry, self).__init__( @@ -39,7 +41,9 @@ def __init__(self, title=None, username=None, password=None, url=None, version=version, expires=expires, expiry_time=expiry_time, - icon=icon + icon=icon, + customicon=customicon, + meta=meta ) self._element.append(E.String(E.Key('Title'), E.Value(title))) self._element.append(E.String(E.Key('UserName'), E.Value(username))) @@ -125,14 +129,6 @@ def notes(self): def notes(self, value): return self._set_string_field('Notes', value) - @property - def icon(self): - return self._get_subelement_text('IconID') - - @icon.setter - def icon(self, value): - return self._set_subelement_text('IconID', value) - @property def tags(self): val = self._get_subelement_text('Tags') @@ -147,7 +143,7 @@ def tags(self, value): @property def history(self): if self._element.find('History') is not None: - return [Entry(element=x, version=self._version) for x in self._element.find('History').findall('Entry')] + return [Entry(element=x, version=self._version, meta=self._meta) for x in self._element.find('History').findall('Entry')] @history.setter def history(self, value): @@ -241,9 +237,9 @@ def __str__(self): def __eq__(self, other): return ( (self.title, self.username, self.password, self.url, - self.notes, self.icon, self.tags, self.atime, self.ctime, + self.notes, self.icon, self.customicon, self.tags, self.atime, self.ctime, self.mtime, self.expires, self.uuid) == (other.title, other.username, other.password, other.url, - other.notes, other.icon, other.tags, other.atime, other.ctime, + other.notes, other.icon, other.customicon, other.tags, other.atime, other.ctime, other.mtime, other.expires, other.uuid) ) diff --git a/pykeepass/group.py b/pykeepass/group.py index a5cb1840..01000eaf 100644 --- a/pykeepass/group.py +++ b/pykeepass/group.py @@ -9,14 +9,14 @@ class Group(BaseElement): - def __init__(self, name=None, element=None, icon=None, notes=None, - version=None, expires=None, expiry_time=None): + def __init__(self, name=None, element=None, icon=None, customicon=None, notes=None, + version=None, expires=None, expiry_time=None, meta=None): assert type(version) is tuple, 'The provided version is not a tuple, but a {}'.format( type(version) ) self._version = version - + self._meta = meta if element is None: super(Group, self).__init__( @@ -24,7 +24,9 @@ def __init__(self, name=None, element=None, icon=None, notes=None, version=version, expires=expires, expiry_time=expiry_time, - icon=icon + icon=icon, + customicon=customicon, + meta=meta ) self._element.append(E.Name(name)) if notes: diff --git a/pykeepass/pykeepass.py b/pykeepass/pykeepass.py index 5e28d6c4..084ba3e7 100644 --- a/pykeepass/pykeepass.py +++ b/pykeepass/pykeepass.py @@ -27,12 +27,6 @@ def __init__(self, filename, password=None, keyfile=None): self.read(password=password, keyfile=keyfile) - def __enter__(self): - return self - - def __exit__(self, typ, value, tb): - del self.kdbx - def read(self, filename=None, password=None, keyfile=None): self.password = password self.keyfile = keyfile @@ -97,6 +91,10 @@ def groups(self): def entries(self): return self.find_entries_by_title('.*', regex=True) + @property + def meta(self): + return self.tree.find('Meta') + def dump_xml(self, outfile): ''' Dump the content of the database to a file @@ -123,9 +121,9 @@ def _xpath(self, xpath_str, tree=None): res = [] for r in result: if r.tag == 'Entry': - res.append(Entry(element=r, version=self.version)) + res.append(Entry(element=r, version=self.version, meta=self.meta)) elif r.tag == 'Group': - res.append(Group(element=r, version=self.version)) + res.append(Group(element=r, version=self.version, meta=self.meta)) else: res.append(r) return res @@ -268,10 +266,12 @@ def find_groups_by_notes(self, notes, regex=False, flags=None, ) # creates a new group and all parent groups, if necessary - def add_group(self, destination_group, group_name, icon=None, notes=None): + def add_group(self, destination_group, group_name, icon=None, customicon=None, notes=None): logger.debug('Creating group {}'.format(group_name)) - if icon: + if customicon: + group = Group(name=group_name, customicon=customicon, notes=notes, version=self.version, meta=self.meta) + elif icon: group = Group(name=group_name, icon=icon, notes=notes, version=self.version) else: group = Group(name=group_name, notes=notes, version=self.version) @@ -409,7 +409,7 @@ def find_entries_by_string(self, string, regex=False, flags=None, def add_entry(self, destination_group, title, username, password, url=None, notes=None, expiry_time=None, - tags=None, icon=None, force_creation=False): + tags=None, icon=None, customicon=None, force_creation=False): entries = self.find_entries( title=title, @@ -437,7 +437,9 @@ def add_entry(self, destination_group, title, username, expires=True if expiry_time else False, expiry_time=expiry_time, icon=icon, - version=self.version + customicon=customicon, + version=self.version, + meta=self.meta ) destination_group.append(entry) diff --git a/tests/tests.py b/tests/tests.py index f22520d9..a1ee9c73 100644 --- a/tests/tests.py +++ b/tests/tests.py @@ -314,6 +314,7 @@ def test_set_and_get_fields(self): entry.expires = False entry.expiry_time = changed_time entry.icon = icons.GLOBE + entry.customicon = "9" entry.set_custom_property('foo', 'bar') self.assertEqual(entry.title, changed_string + 'title') @@ -322,6 +323,7 @@ def test_set_and_get_fields(self): self.assertEqual(entry.url, changed_string + 'url') self.assertEqual(entry.notes, changed_string + 'notes') self.assertEqual(entry.icon, icons.GLOBE) + self.assertEqual(entry.customicon, None) self.assertEqual(entry.get_custom_property('foo'), 'bar') self.assertIn('foo', entry.custom_properties) # test time properties @@ -391,12 +393,6 @@ def test_db_info(self): def tearDown(self): os.remove(os.path.join(base_dir, 'change_creds.kdbx')) -class CtxManagerTests(unittest.TestCase): - def test_ctx_manager(self): - with PyKeePass(os.path.join(base_dir, 'test.kdbx'), password='passw0rd', keyfile=base_dir + '/test.key') as kp: - results = kp.find_entries_by_username('foobar_user', first=True) - self.assertEqual('foobar_user', results.username) - class KDBXTests(unittest.TestCase): def test_open_save(self):