BATCAVE/ffstatus/filestorage.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import cPickle as pickle
import logging
import os
import shutil

from .basestorage import BaseStorage


class FileStorage(BaseStorage):
    """Provides file-based persistency for BaseStorage"""

    __data = None

    def __init__(self, storage_dir):
        self.logger = logging.getLogger('Storage')
        self.storage_dir = storage_dir

        self.logger.debug('Using storage at \'{0}\''.format(storage_dir))
        self.storage_filename = os.path.join(storage_dir, 'storage.dat')

        self.storage_file = None

    def open(self):
        self.storage_file = open(self.storage_filename, 'a+b')
        loaded_data = {}
        try:
            self.storage_file.seek(0, os.SEEK_SET)
            loaded_data = pickle.load(self.storage_file)

        except EOFError:
            self.logger.warn(
                'The storage file was empty.' +
                'Assuming this is the first start with this storage directory.')

        except Exception as err:
            self.logger.error(
                'Loading the storage failed. Please resolve the error ' +
                'and/or remove the storage file \'%s\': %s',
                self.storage_filename,
                str(err),
            )
            raise err

        self.logger.info('Opened storage with %d entries.', len(loaded_data))
        self.__data = loaded_data

    def save(self):
        if self.storage_file is None:
            return

        self.storage_file.seek(0, os.SEEK_SET)
        self.storage_file.truncate()
        pickle.dump(self.__data, self.storage_file, protocol=2)
        self.storage_file.flush()

        # make an auto-backup of the just-written file
        shutil.copyfile(
            self.storage_filename,
            self.storage_filename + '.autobackup')

    def close(self):
        self.save()
        self.storage_file.close()
        self.storage_file = None
        BaseStorage.close(self)

    def get_all_nodes_raw(self):
        """Gets all nodes as dict."""

        result = {}
        for nodeid in self.__data:
            if nodeid.startswith('__'):
                continue
            node = self.__data[nodeid]
            result[nodeid] = node

        return result

    def set_node_data(self, key, data):
        self.__data[key] = data
        self.save()

    def get_vpn_keys(self):
        """Gets a list of VPN keys."""
        if not self.DATAKEY_VPN in self.__data:
            return {}

        vpn = [key for key in self.__data[self.DATAKEY_VPN]]
        return vpn

    def get_vpn_item(self, key, create=False):
        """Get the VPN entry with the specified key."""
        if not self.DATAKEY_VPN in self.__data:
            if not create:
                return None
            self.__data[self.DATAKEY_VPN] = {}

        if not key in self.__data[self.DATAKEY_VPN]:
            if not create:
                return None
            self.__data[self.DATAKEY_VPN][key] = {'active': {}, 'last': {}}

        return self.__data[self.DATAKEY_VPN][key]

    def store_vpn_item(self, key, data):
        if not self.DATAKEY_VPN in self.__data:
            self.__data[self.DATAKEY_VPN] = {}
        self.__data[self.DATAKEY_VPN][key] = data
        self.save()