1# -*- coding: utf-8 -*-
2
3from ..state.prekeystore import PreKeyStore
4from ..state.prekeyrecord import PreKeyRecord
5from ..invalidkeyidexception import InvalidKeyIdException
6
7
8class InMemoryPreKeyStore(PreKeyStore):
9    def __init__(self):
10        self.store = {}
11
12    def loadPreKey(self, preKeyId):
13        if preKeyId not in self.store:
14            raise InvalidKeyIdException("No such prekeyRecord!")
15
16        return PreKeyRecord(serialized=self.store[preKeyId])
17
18    def storePreKey(self, preKeyId, preKeyRecord):
19        self.store[preKeyId] = preKeyRecord.serialize()
20
21    def containsPreKey(self, preKeyId):
22        return preKeyId in self.store
23
24    def removePreKey(self, preKeyId):
25        if preKeyId in self.store:
26            del self.store[preKeyId]
27