Optimise icon- and totp-handling.

convert.py

@@ -46,18 +46,6 @@ def getField(item, designation):
     return None
 
 
-def getTotp(item):
-    secure = item["secureContents"]
-    if "sections" in secure:
-        for section in secure["sections"]:
-            if not "fields" in section:
-                continue
-            for field in section["fields"]:
-                if field["t"] == "totp":
-                    return field["v"]
-    return None
-
-
 ICON_MAP = {
     "112": "GEAR",   # API Credential
     "wallet.financial.BankAccountUS": "DOLLAR_SIGN",   # Bank Account
@@ -83,7 +71,7 @@ ICON_MAP = {
 
 for item in opif:
 
-    # Make sure target group exists
+    # Determine group/folder
     item_type_name = item.type_name
     target_group_name = "{}s".format(item_type_name)   # plural for group
 
@@ -93,19 +81,18 @@ for item in opif:
     # Add entry to KeePass
     entry = kp.add_entry(target_group_name, item["title"])
 
-    # Set icon
+    # Icon
     kp_icon = ICON_MAP[item.type]
-    if kp_icon in dir(pykeepass.icons):
-        kp_icon_id = getattr(pykeepass.icons, kp_icon)
-    else:
-        # FIXME: Assume kp_icon is already ID, needed b/c icon 12 is missing from pykeepass.icons
-        kp_icon_id = kp_icon
-    entry.icon = kp_icon_id
-
-    secure = item["secureContents"]
+    kp.set_icon(kp_icon)
 
     # Tags
-    entry.tags = item.get_tags()
+    kp.set_tags(item.get_tags())
+
+
+
+
+
+    secure = item["secureContents"]
 
     # Username
     if "username" in secure:
@@ -124,10 +111,9 @@ for item in opif:
             entry.password = new_password
 
     # TOTP
-    totp = getTotp(item)
+    totp = item.get_totp()
     if totp:
-        entry.set_custom_property("TimeOtp-Secret-Base32", totp)
-        entry.set_custom_property("otp", "otpauth://totp/Sample:username?secret={}&algorithm=SHA1&digits=6&period=30&issuer=Sample".format(quote_plus(totp)))
+        kp.add_totp(totp)
 
     # Other web fields
     if "fields" in secure:

kpwriter.py

@@ -1,10 +1,12 @@
+import pykeepass.icons
+from urllib.parse import quote_plus
 from pykeepass import create_database
 
 
 class KpWriter:
     def __init__(self, filename, password="test"):
         self.kp = create_database(filename, password)
-        self.last_entry = None
+        self.current_entry = None
 
     def add_entry(self, dest_group_name, title):
         # Find group and create if not yet there
@@ -13,8 +15,26 @@ class KpWriter:
             # TODO: Handle nested groups?
             group = self.kp.add_group(self.kp.root_group, dest_group_name)
 
-        self.last_entry = self.kp.add_entry(group, title, "", "")
-        return self.last_entry
+        self.current_entry = self.kp.add_entry(group, title, "", "")
+        return self.current_entry
 
     def save(self):
         return self.kp.save()
+
+    def set_icon(self, icon_id):
+        if icon_id in dir(pykeepass.icons):
+            kp_icon_id = getattr(pykeepass.icons, icon_id)
+        else:
+            # FIXME: Assume kp_icon is already ID, needed b/c icon 12 is missing from pykeepass.icons
+            kp_icon_id = icon_id
+        self.current_entry.icon = kp_icon_id
+
+    def set_tags(self, tag_list):
+        self.current_entry.tags = tag_list
+
+    def add_totp(self, init_string, otp_url=None):
+        if not otp_url:
+            otp_url = "otpauth://totp/Sample:username?secret={}&algorithm=SHA1&digits=6&period=30&issuer=Sample".format(quote_plus(init_string))
+        # TODO: Support multiple / don't overwrite
+        self.current_entry.set_custom_property("TimeOtp-Secret-Base32", init_string)
+        self.current_entry.set_custom_property("otp", otp_url)

onepif/OnepifEntry.py

@@ -41,6 +41,16 @@ class OnepifEntry():
             return []
         return self.raw["openContents"]["tags"]
 
+    def get_totp(self):
+        if "sections" in self.raw["secureContents"]:
+            for section in self.raw["secureContents"]["sections"]:
+                if "fields" not in section:
+                    continue
+                for field in section["fields"]:
+                    if field["n"][:5] == "TOTP_":
+                        return field["v"]
+        return None
+
     def is_trash(self):
         if "trashed" in self.raw:
             return self.raw["trashed"]