Handle difference in pickle protocol between python 2.x and 3.x

[imap-fix-internaldate] / src / caching_data.py

'''
caching_data.py - The module contains the CachingData class.

Copyright (c) 2012 Intra2net AG
Author: Plamen Dimitrov

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
'''
import os, platform, tempfile
import pickle
import logging
from mailbox_state import MailboxState

CACHE_FILENAME = "message_cache.dat"
CACHE_VERSION = "1"

class CachingData:
    """This class is responsible for the caching of data."""
    
    # class attributes
    # integer for version of the cache
    version = None
    # boolean flag which indicates fallback mode of the cache
    fallback_to_date_header = None
    # dictionary of usernames as keys and dictionaries as values
    # the second dictionaries have unique mailbox keys and mailboxes as values
    data = None

    def __init__(self, fallback_mode):
        # open data file or create one and initialize date if not found
        try:
            cachefile = open(CACHE_FILENAME, 'rb')
            cache_info, self.data = pickle.load(cachefile)
            cache_info = cache_info.split(' ')
            self.version = cache_info[0]
            if(self.version != CACHE_VERSION):
                raise IOError("Cache file has version %s and the script version is %s" % (self.version, CACHE_VERSION))
            self.fallback_to_date_header = cache_info[1]
            if(self.fallback_to_date_header != str(fallback_mode)):
                raise IOError("Cache file date fallback mode setting is different than current settings")
            logging.info("Cache file %s loaded", CACHE_FILENAME)
            logging.info("%s users found.", len(self.data))
        except (IOError, ValueError) as ex:
            logging.warning("Couldn't load cache file %s: %s", CACHE_FILENAME, ex)
            logging.warning("DELETING CACHE")
            self.version = CACHE_VERSION
            stored_cache_info = self.version + ' ' + str(fallback_mode)
            self.data = {}
            with open(CACHE_FILENAME, 'wb') as cachefile:
                pickle.dump((stored_cache_info, self.data), cachefile)

    def __del__(self):
        # create temporary file first
        location = os.path.dirname(CACHE_FILENAME)    
        file_descriptor, tmpname = tempfile.mkstemp(dir=location)
        try:
            cachefile = os.fdopen(file_descriptor, 'wb')

            # prepare data based on a save flag
            saved_data = {}
            for user in self.data:
                saved_data[user] = {}
                for box_key in self.data[user]:
                    if(self.data[user][box_key].needs_save):
                        saved_data[user][box_key] = self.data[user][box_key]
                        logging.debug("The mailbox %s will be saved.", saved_data[user][box_key].name)
                if(len(saved_data[user])==0):
                    del saved_data[user]
                    logging.debug("The user %s will not be saved.", user)
            self.data = saved_data
            # avoid test mode or cases where nothing needs saving
            if(len(saved_data)==0):
                cachefile.close()
                os.unlink(tmpname)
                return

            # serialize in file
            stored_cache_info = self.version + ' ' + self.fallback_to_date_header
            pickle.dump((stored_cache_info, self.data), cachefile)
            logging.debug("%s users stored.", len(self.data))

            # handle windows non-atomic rename
            if(platform.system()=='Windows'):
                if(os.path.exists(CACHE_FILENAME)):
                    cachefile.close()
                    os.unlink(CACHE_FILENAME)

            os.rename(tmpname, CACHE_FILENAME)
        except:
            # clean up temporary file
            os.unlink(tmpname)
            raise

        logging.info("Wrote cache file %s", CACHE_FILENAME)

    def retrieve_cached_mailbox(self, name, uidvalidity, user):
        """Retrieve a cached mailbox or create it."""
        box_key = name.strip('"') + uidvalidity
        if(user not in self.data):
            self.data[user] = {}
            logging.debug("New user %s cached.", user)
        if(box_key not in self.data[user]):
            self.data[user][box_key] = MailboxState(name, uidvalidity, user)
            logging.debug("New mailbox %s cached.", box_key)
        return self.data[user][box_key]
    
    def report_conflicts(self):
        """Write a date conflicts report in a file."""
        with open("conflict_stats.txt", 'w') as statsfile:
            owner_total_conflicts = {}
            owner_total_missing = {}
            for user in self.data:
                statsfile.write("user: %s\n" % user)
                owner_total_conflicts[user] = 0
                owner_total_missing[user] = 0
                for box_key in self.data[user]:
                    owner_total_conflicts[user] += self.data[user][box_key].date_conflicts
                    owner_total_missing[user] += self.data[user][box_key].no_received_field
                    statsfile.write("date conflicts: %-15.15s missing header: %-15.15s mailbox: %s\n"\
                                    % (self.data[user][box_key].date_conflicts,
                                            self.data[user][box_key].no_received_field,
                                            self.data[user][box_key].name))
                statsfile.write("date conflicts: %-15.15s missing header: %-15.15s TOTAL \n\n"\
                                % (owner_total_conflicts[user], owner_total_missing[user]))
        return