customicon support

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

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')

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,20 +27,23 @@ 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__(
                 element=Element('Entry'),
                 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)
         )

pykeepass/group.py

@@ -9,22 +9,24 @@
 
 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__(
                 element=Element('Group'),
                 version=version,
                 expires=expires,
                 expiry_time=expiry_time,
-                icon=icon
+                icon=icon,
+                customicon=customicon,
+                meta=meta
             )
             self._element.append(E.Name(name))
             if notes:

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)
 

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):