#!/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_node(self, id): """Gets the node by its id, if any.""" return self.__data.get(id) 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()