add check for none in part_map special case

1 files changed, 7 insertions, 6 deletions

diff --git a/mail/src/leap/mail/walk.py b/mail/src/leap/mail/walk.py
index 856daa3..30cb70a 100644
--- a/mail/src/leap/mail/walk.py
+++ b/mail/src/leap/mail/walk.py
@@ -151,12 +151,13 @@ def walk_msg_tree(parts, body_phash=None):
 
     if all(x == 1 for x in pv):
         # special case in the rightmost element
-        main_pmap = parts[0]['part_map']
-        last_part = max(main_pmap.keys())
-        main_pmap[last_part]['part_map'] = {}
-        for partind in range(len(pv) - 1):
-            print partind+1, len(parts)
-            main_pmap[last_part]['part_map'][partind] = parts[partind+1]
+        main_pmap = parts[0].get(PART_MAP, None)
+        if main_pmap is not None:
+            last_part = max(main_pmap.keys())
+            main_pmap[last_part][PART_MAP] = {}
+            for partind in range(len(pv) - 1):
+                print partind+1, len(parts)
+                main_pmap[last_part][PART_MAP][partind] = parts[partind + 1]
 
     outer = parts[0]
     outer.pop(HEADERS)