filestorage.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115
  1. #!/usr/bin/python
  2. # -*- coding: utf-8 -*-
  3. import cPickle as pickle
  4. import logging
  5. import os
  6. import shutil
  7. from .basestorage import BaseStorage
  8. class FileStorage(BaseStorage):
  9. """Provides file-based persistency for BaseStorage"""
  10. DATAKEY_VPN = '__VPN__'
  11. FIELDKEY_UPDATED = '__UPDATED__'
  12. __data = None
  13. def __init__(self, storage_dir):
  14. self.logger = logging.getLogger('Storage')
  15. self.storage_dir = storage_dir
  16. self.logger.debug('Using storage at \'{0}\''.format(storage_dir))
  17. self.storage_filename = os.path.join(storage_dir, 'storage.dat')
  18. self.storage_file = None
  19. def open(self):
  20. self.storage_file = open(self.storage_filename, 'a+b')
  21. loaded_data = {}
  22. try:
  23. self.storage_file.seek(0, os.SEEK_SET)
  24. loaded_data = pickle.load(self.storage_file)
  25. except EOFError:
  26. self.logger.warn(
  27. 'The storage file was empty.' +
  28. 'Assuming this is the first start with this storage directory.')
  29. except Exception as err:
  30. self.logger.error(
  31. 'Loading the storage failed. Please resolve the error ' +
  32. 'and/or remove the storage file \'%s\': %s',
  33. self.storage_filename,
  34. str(err),
  35. )
  36. raise err
  37. self.logger.info('Opened storage with %d entries.', len(loaded_data))
  38. self.__data = loaded_data
  39. def save(self):
  40. if self.storage_file is None:
  41. return
  42. self.storage_file.seek(0, os.SEEK_SET)
  43. self.storage_file.truncate()
  44. pickle.dump(self.__data, self.storage_file, protocol=2)
  45. self.storage_file.flush()
  46. # make an auto-backup of the just-written file
  47. shutil.copyfile(
  48. self.storage_filename,
  49. self.storage_filename + '.autobackup')
  50. def close(self):
  51. self.save()
  52. self.storage_file.close()
  53. self.storage_file = None
  54. BaseStorage.close(self)
  55. def get_all_nodes_raw(self):
  56. """Gets all nodes as dict."""
  57. result = {}
  58. for nodeid in self.__data:
  59. if nodeid.startswith('__'):
  60. continue
  61. node = self.__data[nodeid]
  62. result[nodeid] = node
  63. return result
  64. def set_node_data(self, key, data):
  65. self.__data[key] = data
  66. self.save()
  67. def get_vpn_keys(self):
  68. """Gets a list of VPN keys."""
  69. if not self.DATAKEY_VPN in self.__data:
  70. return {}
  71. vpn = [key for key in self.__data[self.DATAKEY_VPN]]
  72. return vpn
  73. def get_vpn_item(self, key, create=False):
  74. """Get the VPN entry with the specified key."""
  75. if not self.DATAKEY_VPN in self.__data:
  76. if not create:
  77. return None
  78. self.__data[self.DATAKEY_VPN] = {}
  79. if not key in self.__data[self.DATAKEY_VPN]:
  80. if not create:
  81. return None
  82. self.__data[self.DATAKEY_VPN][key] = {'active': {}, 'last': {}}
  83. return self.__data[self.DATAKEY_VPN][key]
  84. def store_vpn_item(self, key, data):
  85. if not self.DATAKEY_VPN in self.__data:
  86. self.__data[self.DATAKEY_VPN] = {}
  87. self.__data[self.DATAKEY_VPN][key] = data
  88. self.save()