Make the script more robust.

1 files changed, 59 insertions, 32 deletions

diff --git a/archive_folder.py b/archive_folder.py
index 86afaa0..7c1463f 100755
--- a/archive_folder.py
+++ b/archive_folder.py
@@ -16,7 +16,7 @@ import sys
 from datetime import date, timedelta
 
 logging.basicConfig(format='%(levelname)s:%(funcName)s:%(message)s',
-                    stream=sys.stdout, level=logging.DEBUG)
+                    stream=sys.stdout, level=logging.ERROR)
 log = logging.getLogger('imapArch')
 
 class ServerError(IOError):
@@ -40,7 +40,12 @@ class Message(object):
         self.box = self.folder.box
         self.uid = uid
         self.msg = self.__get_body()
-        self.date = email.utils.parsedate_to_datetime(self.msg['Date'])
+        date_header = self.msg['Date']
+        try:
+            self.date = email.utils.parsedate_to_datetime(date_header)
+        except TypeError:
+            log.warning('Cannot parse date string %s', date_header)
+            self.date = None
         self.subject = self.__get_subject()
 
     def __get_body(self):
@@ -53,11 +58,21 @@ class Message(object):
 
     def __get_subject(self):
         out_str = ''
-        for raw_str, enc in email.header.decode_header(self.msg['Subject']):
-            if isinstance(raw_str, bytes):
-                out_str += raw_str.decode(enc or 'utf8')
-            else:
-                out_str += raw_str
+        subj_str = self.msg['Subject']
+        if subj_str is not None:
+            for raw_str, enc in email.header.decode_header(subj_str):
+                if enc == 'windows-874':
+                    enc = 'iso8859_11'
+                elif enc is None or ('unknown' in enc):
+                    enc = None
+                if isinstance(raw_str, bytes):
+                    try:
+                        out_str += raw_str.decode(enc or 'utf8', 'replace')
+                    except LookupError:
+                        log.error('Cannot find encoding %s', enc)
+                        raise
+                else:
+                    out_str += raw_str
         return out_str
 
     def __str__(self):
@@ -66,16 +81,17 @@ class Message(object):
 
 
 class Folder(object):
-    def __init__(self, box, fld_name, create=False):
-        self.box = box
+    def __init__(self, server, fld_name, create=False):
+        self.server = server
+        self.box = server.box
         self.fld_name = fld_name
         self.selected = False
         self.__create_missing = create
         self.folder_sep = self.__get_separator()
 
-        log.debug('self.__create_missing = %s, self.__list_folder = %s',
-                  self.__create_missing, self.__list_folder())
-        if self.__create_missing and self.__list_folder() is None:
+        log.debug('self.__create_missing = %s, fld_name = %s, exists %s',
+                  self.__create_missing, fld_name, fld_name not in self.server.all_folders)
+        if self.__create_missing and (fld_name not in self.server.all_folders):
             self.__create_folder()
 
     def select(self):
@@ -94,16 +110,20 @@ class Folder(object):
             log.debug('part = %s', part)
             target += '{}{}'.format(part, self.folder_sep)
             fld = target.strip(self.folder_sep)
-            if self.__list_folder(fld) is None:
+            log.debug('fld = %s', fld)
+            if fld not in self.server.all_folders:
+                log.debug('Creating folder %s', fld)
                 self.box.create(fld)
+                self.server.all_folders.append(fld)
                 self.box.subscribe(fld)
 
     def __list_folder(self, name=None, wildcards='*'):
         if name is None:
-            name = self.fld_name
+            name = '{}/'.format(self.fld_name)
         ok, data = self.box.list(name, wildcards)
         if ok != 'OK':
             raise ServerError('Cannot list folder %s' % self.fld_name)
+        log.debug('data = %s', data)
         return data[0]
 
     def __get_separator(self):
@@ -130,7 +150,8 @@ class Folder(object):
                                (search_type.lower(), date_str))
 
         msgs = []
-        for uid in res[0].decode().split(' '):
+        uids = res[0].decode().split()
+        for uid in uids:
             msgs.append(Message(self, uid))
         return msgs
 
@@ -190,7 +211,7 @@ class EmailServer(object):
             else config['general']['account']
         self.cfg = dict(config.items(acc_name))
         self.archive_root = archive_root
-        self.box = self.__login(**self.cfg)
+        self.box, self.all_folders = self.__login(**self.cfg)
         # self.box.debug = 4
 
     def __login(self, host='localhost', username=None, password=None, ssl=None):
@@ -204,7 +225,14 @@ class EmailServer(object):
         capas = data[0].decode()
         log.debug('capas = %s', capas)
         box.features_present = Capas._make(['MOVE' in capas, 'UIDPLUS' in capas])
-        return box
+
+        ok, data = box.list()
+        if ok != 'OK':
+            raise IOError('Cannot list folders!')
+
+        folders = [FOLDER_RE.split(x.decode())[1] for x in data]
+
+        return box, folders
 
     @staticmethod
     def get_config():
@@ -222,36 +250,35 @@ class EmailServer(object):
         :param: folder_name
         :param: before_date
         """
+        print('Archiving {}'.format(folder_name))
         copy_cache = {}
-        fld = Folder(self.box, folder_name).select()
+        fld = Folder(self, folder_name).select()
         before_str = before_date.strftime('%d-%b-%Y')
         for msg in fld.emails_before(before_str):
-            arch_folder = fld.get_archive_folder(msg, self.archive_root)
-            if arch_folder in copy_cache:
-                copy_cache[arch_folder].append(msg)
-            else:
-                copy_cache[arch_folder] = [msg]
+            sys.stdout.write('.')
+            if msg.date is not None:
+                arch_folder = fld.get_archive_folder(msg, self.archive_root)
+                if arch_folder in copy_cache:
+                    copy_cache[arch_folder].append(msg)
+                else:
+                    copy_cache[arch_folder] = [msg]
+        sys.stdout.write('\n')
 
         for key in copy_cache:
-            log.info('***** %s:', key)
+            log.debug('***** %s:', key)
             for msg in copy_cache[key]:
-                log.info('\tmsg: %s', str(msg))
+                log.debug('\tmsg: %s', str(msg))
 
         # Go through the cache and make moves.
         for arch_dir in copy_cache:
             msg_ids = ','.join([x.uid for x in copy_cache[arch_dir]])
             log.debug('arch_dir = %s, msgs = %s', arch_dir,
                       msg_ids)
-            dir = Folder(self.box, arch_dir, create=True)
+            dir = Folder(self, arch_dir, create=True)
             dir.move_messages(msg_ids)
 
     def list_folders(self):
-        ok, data = self.box.list()
-        if ok != 'OK':
-            raise IOError('Cannot list folders!')
-
-        return sorted([FOLDER_RE.split(x.decode())[1] for x in data],
-                      key=str.lower)
+        return sorted(self.all_folders, key=str.lower)
 
     def __enter__(self):
         return self