PyORAM/src/pyoram/encrypted_storage/encrypted_block_storage.py

   1 __all__ = ('EncryptedBlockStorage',)
   2
   3 import struct
   4 import hmac
   5 import hashlib
   6
   7 from pyoram.storage.block_storage import (BlockStorageInterface,
   8                                           BlockStorageTypeFactory)
   9 from pyoram.crypto.aes import AES
  10
  11 import six
  12
  13 class EncryptedBlockStorageInterface(BlockStorageInterface):
  14
  15     #
  16     # Abstract Interface
  17     #
  18
  19     @property
  20     def key(self, *args, **kwds):
  21         raise NotImplementedError                      # pragma: no cover
  22     @property
  23     def raw_storage(self, *args, **kwds):
  24         raise NotImplementedError                      # pragma: no cover
  25
  26 class EncryptedBlockStorage(EncryptedBlockStorageInterface):
  27
  28     _index_struct_string = "!"+("x"*hashlib.sha384().digest_size)+"?"
  29     _index_offset = struct.calcsize(_index_struct_string)
  30     _verify_struct_string = "!LLL"
  31     _verify_size = struct.calcsize(_verify_struct_string)
  32
  33     def __init__(self, storage, **kwds):
  34         self._key = kwds.pop('key', None)
  35         if self._key is None:
  36             raise ValueError(
  37                 "An encryption key is required using "
  38                 "the 'key' keyword.")
  39         if isinstance(storage, BlockStorageInterface):
  40             storage_owned = False
  41             self._storage = storage
  42             if len(kwds):
  43                 raise ValueError(
  44                     "Keywords not used when initializing "
  45                     "with a storage device: %s"
  46                     % (str(kwds)))
  47         else:
  48             storage_owned = True
  49             storage_type = kwds.pop('storage_type', 'file')
  50             self._storage = \
  51                 BlockStorageTypeFactory(storage_type)(storage, **kwds)
  52
  53         try:
  54             header_data = AES.GCMDec(self._key,
  55                                      self._storage.header_data)
  56             (self._ismodegcm,) = struct.unpack(
  57                 self._index_struct_string,
  58                 header_data[:self._index_offset])
  59             self._verify_digest = header_data[:hashlib.sha384().digest_size]
  60
  61             verify = hmac.HMAC(
  62                 key=self.key,
  63                 msg=struct.pack(self._verify_struct_string,
  64                                 self._storage.block_size,
  65                                 self._storage.block_count,
  66                                 len(self._storage.header_data)),
  67                 digestmod=hashlib.sha384)
  68             if verify.digest() != self._verify_digest:
  69                 raise ValueError(
  70                     "HMAC of plaintext index data does not match")
  71             if self._ismodegcm:
  72                 self._encrypt_block_func = AES.GCMEnc
  73                 self._decrypt_block_func = AES.GCMDec
  74             else:
  75                 self._encrypt_block_func = AES.CTREnc
  76                 self._decrypt_block_func = AES.CTRDec
  77         except:
  78             if storage_owned:
  79                 self._storage.close()
  80             raise
  81
  82     #
  83     # Define EncryptedBlockStorageInterface Methods
  84     #
  85
  86     @property
  87     def key(self):
  88         return self._key
  89
  90     @property
  91     def raw_storage(self):
  92         return self._storage
  93
  94     #
  95     # Define BlockStorageInterface Methods
  96     #
  97
  98     def clone_device(self):
  99         return EncryptedBlockStorage(self._storage.clone_device(),
 100                                      key=self.key)
 101
 102     @classmethod
 103     def compute_storage_size(cls,
 104                              block_size,
 105                              block_count,
 106                              aes_mode='ctr',
 107                              storage_type='file',
 108                              ignore_header=False,
 109                              **kwds):
 110         assert (block_size > 0) and (block_size == int(block_size))
 111         assert (block_count > 0) and (block_count == int(block_count))
 112         assert aes_mode in ('ctr', 'gcm')
 113         if not isinstance(storage_type, BlockStorageInterface):
 114             storage_type = BlockStorageTypeFactory(storage_type)
 115
 116         if aes_mode == 'ctr':
 117             extra_block_data = AES.block_size
 118         else:
 119             assert aes_mode == 'gcm'
 120             extra_block_data = 2 * AES.block_size
 121         if ignore_header:
 122             return (extra_block_data * block_count) + \
 123                     storage_type.compute_storage_size(
 124                         block_size,
 125                         block_count,
 126                         ignore_header=True,
 127                         **kwds)
 128         else:
 129             return cls._index_offset + \
 130                    2 * AES.block_size + \
 131                    (extra_block_data * block_count) + \
 132                    storage_type.compute_storage_size(
 133                        block_size,
 134                        block_count,
 135                        ignore_header=False,
 136                        **kwds)
 137
 138     @classmethod
 139     def setup(cls,
 140               storage_name,
 141               block_size,
 142               block_count,
 143               aes_mode='ctr',
 144               key_size=None,
 145               key=None,
 146               storage_type='file',
 147               initialize=None,
 148               **kwds):
 149
 150         if (key is not None) and (key_size is not None):
 151             raise ValueError(
 152                 "Only one of 'key' or 'keysize' keywords can "
 153                 "be specified at a time")
 154         if key is None:
 155             if key_size is None:
 156                 key_size = AES.key_sizes[-1]
 157             if key_size not in AES.key_sizes:
 158                 raise ValueError(
 159                     "Invalid key size: %s" % (key_size))
 160             key = AES.KeyGen(key_size)
 161         else:
 162             if len(key) not in AES.key_sizes:
 163                 raise ValueError(
 164                     "Invalid key size: %s" % (len(key)))
 165
 166         if (block_size <= 0) or (block_size != int(block_size)):
 167             raise ValueError(
 168                 "Block size (bytes) must be a positive integer: %s"
 169                 % (block_size))
 170
 171         ismodegcm = None
 172         encrypt_block_func = None
 173         encrypted_block_size = block_size
 174         if aes_mode == 'ctr':
 175             ismodegcm = False
 176             encrypt_block_func = AES.CTREnc
 177             encrypted_block_size += AES.block_size
 178         elif aes_mode == 'gcm':
 179             ismodegcm = True
 180             encrypt_block_func = AES.GCMEnc
 181             encrypted_block_size += (2 * AES.block_size)
 182         else:
 183             raise ValueError(
 184                 "AES encryption mode must be one of 'ctr' or 'gcm'. "
 185                 "Invalid value: %s" % (aes_mode))
 186         assert ismodegcm is not None
 187         assert encrypt_block_func is not None
 188
 189         if not isinstance(storage_type, BlockStorageInterface):
 190             storage_type = BlockStorageTypeFactory(storage_type)
 191
 192         if initialize is None:
 193             zeros = bytes(bytearray(block_size))
 194             initialize = lambda i: zeros
 195         def encrypted_initialize(i):
 196             return encrypt_block_func(key, initialize(i))
 197         kwds['initialize'] = encrypted_initialize
 198
 199         user_header_data = kwds.get('header_data', bytes())
 200         if type(user_header_data) is not bytes:
 201             raise TypeError(
 202                 "'header_data' must be of type bytes. "
 203                 "Invalid type: %s" % (type(user_header_data)))
 204         # we generate the first time simply to
 205         # compute the length
 206         tmp = hmac.HMAC(
 207             key=key,
 208             msg=struct.pack(cls._verify_struct_string,
 209                             encrypted_block_size,
 210                             block_count,
 211                             0),
 212             digestmod=hashlib.sha384).digest()
 213         header_data = bytearray(struct.pack(cls._index_struct_string,
 214                                             ismodegcm))
 215         header_data[:hashlib.sha384().digest_size] = tmp
 216         header_data = header_data + user_header_data
 217         header_data = AES.GCMEnc(key, bytes(header_data))
 218         # now that we know the length of the header data
 219         # being sent to the underlying storage we can
 220         # compute the real hmac
 221         verify_digest = hmac.HMAC(
 222             key=key,
 223             msg=struct.pack(cls._verify_struct_string,
 224                             encrypted_block_size,
 225                             block_count,
 226                             len(header_data)),
 227             digestmod=hashlib.sha384).digest()
 228         header_data = bytearray(struct.pack(cls._index_struct_string,
 229                                             ismodegcm))
 230         header_data[:hashlib.sha384().digest_size] = verify_digest
 231         header_data = header_data + user_header_data
 232         kwds['header_data'] = AES.GCMEnc(key, bytes(header_data))
 233
 234         return EncryptedBlockStorage(
 235             storage_type.setup(storage_name,
 236                                encrypted_block_size,
 237                                block_count,
 238                                **kwds),
 239             key=key)
 240
 241     @property
 242     def header_data(self):
 243         return AES.GCMDec(self._key,
 244                           self._storage.header_data)\
 245                           [self._index_offset:]
 246
 247     @property
 248     def block_count(self):
 249         return self._storage.block_count
 250
 251     @property
 252     def block_size(self):
 253         if self._ismodegcm:
 254             return self._storage.block_size - 2 * AES.block_size
 255         else:
 256             return self._storage.block_size - AES.block_size
 257
 258     @property
 259     def storage_name(self):
 260         return self._storage.storage_name
 261
 262     def update_header_data(self, new_header_data):
 263         self._storage.update_header_data(
 264             AES.GCMEnc(
 265                 self.key,
 266                 AES.GCMDec(self._key,
 267                            self._storage.header_data)\
 268                            [:self._index_offset] + \
 269                            new_header_data))
 270
 271     def close(self):
 272         self._storage.close()
 273
 274
 275
 276
 277
 278
 279
 280
 281
 282
 283
 284
 285     def read_block(self, i):
 286         a = self._storage.read_block(i)
 287         return self._decrypt_block_func(self._key,a)
 288
 289     def read_blocks(self, indices, *args, **kwds):
 290         a = self._storage.read_blocks(indices, *args, **kwds)
 291         return [self._decrypt_block_func(self._key, b) for b in a]
 292
 293     def yield_blocks(self, indices, *args, **kwds):
 294         for b in self._storage.yield_blocks(indices, *args, **kwds):
 295             yield self._decrypt_block_func(self._key, b)
 296
 297
 298     def write_block(self, i, block, *args, **kwds):
 299         a = self._encrypt_block_func(self._key, block)
 300         self._storage.write_block(i, a,*args, **kwds)
 301
 302     def write_blocks(self, indices, blocks, *args, **kwds):
 303         enc_blocks = []
 304         for i, b in zip(indices, blocks):
 305             enc_blocks.append(self._encrypt_block_func(self._key, b))
 306
 307
 308         self._storage.write_blocks(indices, enc_blocks, *args, **kwds)
 309
 310
 311
 312
 313
 314
 315
 316
 317
 318
 319
 320     @property
 321     def bytes_sent(self):
 322         return self._storage.bytes_sent
 323
 324     @property
 325     def bytes_received(self):
 326         return self._storage.bytes_received