class Curve:
    DJB_TYPE = 5
    # always DJB curve25519 keys

    @staticmethod
    def generatePrivateKey():
        rand = os.urandom(32)
        return _curve.generatePrivateKey(rand)

    @staticmethod
    def generatePublicKey(privateKey):
        return _curve.generatePublicKey(privateKey)

    @staticmethod
    def generateKeyPair():
        from .djbec import DjbECPublicKey, DjbECPrivateKey

        privateKey = Curve.generatePrivateKey()
        publicKey = Curve.generatePublicKey(privateKey)
        return ECKeyPair(DjbECPublicKey(publicKey), DjbECPrivateKey(privateKey))

    @staticmethod
    def decodePoint(_bytes, offset=0):
        type = _bytes[0]  # byte appears to be automatically converted to an integer??

        if type == Curve.DJB_TYPE:
            from .djbec import DjbECPublicKey

            type = _bytes[offset] & 0xFF
            if type != Curve.DJB_TYPE:
                raise InvalidKeyException("Unknown key type: %s " % type)
            keyBytes = _bytes[offset+1:][:32]
            return DjbECPublicKey(bytes(keyBytes))
        else:
            raise InvalidKeyException("Unknown key type: %s" % type)

    @staticmethod
    def decodePrivatePoint(_bytes):
        from .djbec import DjbECPrivateKey
        return DjbECPrivateKey(bytes(_bytes))

    @staticmethod
    def calculateAgreement(publicKey, privateKey):
        """
        :type publicKey: ECPublicKey
        :type privateKey: ECPrivateKey
        """
        if publicKey.getType() != privateKey.getType():
            raise InvalidKeyException("Public and private keys must be of the same type!")

        if publicKey.getType() == Curve.DJB_TYPE:
            return _curve.calculateAgreement(privateKey.getPrivateKey(), publicKey.getPublicKey())
        else:
            raise InvalidKeyException("Unknown type: %s" % publicKey.getType())

    @staticmethod
    def verifySignature(ecPublicSigningKey, message, signature):
        """
        :type ecPublicSigningKey: ECPublicKey
        :type message: bytearray
        :type signature: bytearray
        """

        if ecPublicSigningKey.getType() == Curve.DJB_TYPE:
            result = _curve.verifySignature(ecPublicSigningKey.getPublicKey(), message, signature)
            return result == 0
        else:
            raise InvalidKeyException("Unknown type: %s" % ecPublicSigningKey.getType())

    @staticmethod
    def calculateSignature(privateSigningKey, message):
        """
        :type privateSigningKey: ECPrivateKey
        :type  message: bytearray
        """
        if privateSigningKey.getType() == Curve.DJB_TYPE:
            rand = os.urandom(64)
            res = _curve.calculateSignature(rand, privateSigningKey.getPrivateKey(), message)
            return res
        else:
            raise InvalidKeyException("Unknown type: %s" % privateSigningKey.getType()) class AliceAxolotlParameters:
    def __init__(self, ourIdentityKey, ourBaseKey, theirIdentityKey, theirSignedPreKey,
                 theirRatchetKey, theirOneTimePreKey):
        """
        :type ourBaseKey: ECKeyPair
        :type theirSignedPreKey: ECKeyPair
        :type ourIdentityKey: IdentityKeyPair
        :type theirOneTimePreKey: ECPublicKey
        :type theirRatchetKey: ECKeyPair
        :type theirIdentityKey: IdentityKey
        """
        self.ourBaseKey = ourBaseKey
        self.ourIdentityKey = ourIdentityKey
        self.theirSignedPreKey = theirSignedPreKey
        self.theirRatchetKey = theirRatchetKey
        self.theirIdentityKey = theirIdentityKey
        self.theirOneTimePreKey = theirOneTimePreKey

        if ourBaseKey is None or ourIdentityKey is None or theirSignedPreKey is None \
                or theirRatchetKey is None or theirIdentityKey is None or theirSignedPreKey is None:
            raise ValueError("Null value!")

    def getOurIdentityKey(self):
        return self.ourIdentityKey

    def getOurBaseKey(self):
        return self.ourBaseKey

    def getTheirIdentityKey(self):
        return self.theirIdentityKey

    def getTheirSignedPreKey(self):
        return self.theirSignedPreKey

    def getTheirOneTimePreKey(self):
        return self.theirOneTimePreKey

    def getTheirRatchetKey(self):
        return self.theirRatchetKey

    @staticmethod
    def newBuilder():
        return AliceAxolotlParameters.Builder()

    class Builder:
        def __init__(self):
            self.ourIdentityKey = None
            self.ourBaseKey = None
            self.theirIdentityKey = None
            self.theirSignedPreKey = None
            self.theirRatchetKey = None
            self.theirOneTimePreKey = None

        def setOurIdentityKey(self, ourIdentityKey):
            self.ourIdentityKey = ourIdentityKey
            return self

        def setOurBaseKey(self, ourBaseKey):
            self.ourBaseKey = ourBaseKey
            return self

        def setTheirRatchetKey(self, theirRatchetKey):
            self.theirRatchetKey = theirRatchetKey
            return self

        def setTheirIdentityKey(self, theirIdentityKey):
            self.theirIdentityKey = theirIdentityKey
            return self

        def setTheirSignedPreKey(self, theirSignedPreKey):
            self.theirSignedPreKey = theirSignedPreKey
            return self

        def setTheirOneTimePreKey(self, theirOneTimePreKey):
            self.theirOneTimePreKey = theirOneTimePreKey
            return self

        def create(self):
            return AliceAxolotlParameters(self.ourIdentityKey, self.ourBaseKey, self.theirIdentityKey,
                                          self.theirSignedPreKey, self.theirRatchetKey, self.theirOneTimePreKey) class BobAxolotlParameters:
    def __init__(self, ourIdentityKey, ourSignedPreKey, ourRatchetKey, ourOneTimePreKey,
                 theirIdentityKey, theirBaseKey):
        """
        :type ourIdentityKey: IdentityKeyPair
        :type ourSignedPreKey: ECKeyPair
        :type ourRatchetKey: ECKeyPair
        :type ourOneTimePreKey: ECKeyPair
        :type theirIdentityKey: IdentityKey
        :type theirBaseKey: ECPublicKey
        """
        self.ourIdentityKey = ourIdentityKey
        self.ourSignedPreKey = ourSignedPreKey
        self.ourRatchetKey = ourRatchetKey
        self.ourOneTimePreKey = ourOneTimePreKey
        self.theirIdentityKey = theirIdentityKey
        self.theirBaseKey = theirBaseKey

        if ourIdentityKey is None or ourSignedPreKey is None or ourRatchetKey is None \
                or theirIdentityKey is None or theirBaseKey is None:
            raise ValueError("Null value!")

    def getOurIdentityKey(self):
        return self.ourIdentityKey

    def getOurSignedPreKey(self):
        return self.ourSignedPreKey

    def getOurOneTimePreKey(self):
        return self.ourOneTimePreKey

    def getTheirIdentityKey(self):
        return self.theirIdentityKey

    def getTheirBaseKey(self):
        return self.theirBaseKey

    def getOurRatchetKey(self):
        return self.ourRatchetKey

    @staticmethod
    def newBuilder():
        return BobAxolotlParameters.Builder()

    class Builder:
        def __init__(self):
            self.ourIdentityKey = None
            self.ourSignedPreKey = None
            self.ourOneTimePreKey = None
            self.ourRatchetKey = None
            self.theirIdentityKey = None
            self.theirBaseKey = None

        def setOurIdentityKey(self, ourIdentityKey):
            self.ourIdentityKey = ourIdentityKey
            return self

        def setOurSignedPreKey(self, ourSignedPreKey):
            self.ourSignedPreKey = ourSignedPreKey
            return self

        def setOurOneTimePreKey(self, ourOneTimePreKey):
            self.ourOneTimePreKey = ourOneTimePreKey
            return self

        def setOurRatchetKey(self, ourRatchetKey):
            self.ourRatchetKey = ourRatchetKey
            return self

        def setTheirIdentityKey(self, theirIdentityKey):
            self.theirIdentityKey = theirIdentityKey
            return self

        def setTheirBaseKey(self, theirBaseKey):
            self.theirBaseKey = theirBaseKey
            return self

        def create(self):
            return BobAxolotlParameters(self.ourIdentityKey, self.ourSignedPreKey, self.ourRatchetKey,
                                        self.ourOneTimePreKey, self.theirIdentityKey, self.theirBaseKey) class SymmetricAxolotlParameters:
    def __init__(self, ourBaseKey, ourRatchetKey, ourIdentityKey, theirBaseKey,
                 theirRatchetKey, theirIdentityKey):
        """
        :type ourBaseKey: ECKeyPair
        :type ourRatchetKey: ECKeyPair
        :type ourIdentityKey: IdentityKeyPair
        :type theirBaseKey: ECPublicKey
        :type theirRatchetKey: ECKeyPair
        :type theirIdentityKey: IdentityKey
        """
        self.ourBaseKey = ourBaseKey
        self.ourIdentityKey = ourIdentityKey
        self.ourRatchetKey = ourRatchetKey
        self.theirRatchetKey = theirRatchetKey
        self.theirIdentityKey = theirIdentityKey
        self.theirBaseKey = theirBaseKey

        if ourBaseKey is None or ourIdentityKey is None or ourRatchetKey is None \
                or theirRatchetKey is None or theirIdentityKey is None or theirBaseKey is None:
            raise ValueError("Null value!")

    def getOurBaseKey(self):
        return self.ourBaseKey

    def getOurIdentityKey(self):
        return self.ourIdentityKey

    def getTheirRatchetKey(self):
        return self.theirRatchetKey

    def getTheirIdentityKey(self):
        return self.theirIdentityKey

    def getTheirBaseKey(self):
        return self.theirBaseKey

    def getOurRatchetKey(self):
        return self.ourRatchetKey

    @staticmethod
    def newBuilder():
        return SymmetricAxolotlParameters.Builder()

    class Builder:
        def __init__(self):
            self.ourIdentityKey = None
            self.ourBaseKey = None
            self.ourRatchetKey = None
            self.theirRatchetKey = None
            self.theirIdentityKey = None
            self.theirBaseKey = None

        def setOurIdentityKey(self, ourIdentityKey):
            self.ourIdentityKey = ourIdentityKey
            return self

        def setOurBaseKey(self, ourBaseKey):
            self.ourBaseKey = ourBaseKey
            return self

        def setOurRatchetKey(self, ourRatchetKey):
            self.ourRatchetKey = ourRatchetKey
            return self

        def setTheirRatchetKey(self, theirRatchetKey):
            self.theirRatchetKey = theirRatchetKey
            return self

        def setTheirIdentityKey(self, theirIdentityKey):
            self.theirIdentityKey = theirIdentityKey
            return self

        def setTheirBaseKey(self, theirBaseKey):
            self.theirBaseKey = theirBaseKey
            return self

        def create(self):
            return SymmetricAxolotlParameters(self.ourBaseKey, self.ourRatchetKey, self.ourIdentityKey,
                                              self.theirBaseKey, self.theirRatchetKey, self.theirIdentityKey) class RatchetingSession:
    @staticmethod
    def initializeSession(sessionState, parameters):
        """
        :type sessionState: SessionState
        :type parameters: SymmetricAxolotlParameters
        """
        if RatchetingSession.isAlice(parameters.getOurBaseKey().getPublicKey(), parameters.getTheirBaseKey()):
            aliceParameters = AliceAxolotlParameters.newBuilder()
            aliceParameters.setOurBaseKey(parameters.getOurBaseKey()) \
                .setOurIdentityKey(parameters.getOurIdentityKey()) \
                .setTheirRatchetKey(parameters.getTheirRatchetKey()) \
                .setTheirIdentityKey(parameters.getTheirIdentityKey()) \
                .setTheirSignedPreKey(parameters.getTheirBaseKey()) \
                .setTheirOneTimePreKey(None)
            RatchetingSession.initializeSessionAsAlice(sessionState, aliceParameters.create())
        else:
            bobParameters = BobAxolotlParameters.newBuilder()
            bobParameters.setOurIdentityKey(parameters.getOurIdentityKey()) \
                .setOurRatchetKey(parameters.getOurRatchetKey()) \
                .setOurSignedPreKey(parameters.getOurBaseKey()) \
                .setOurOneTimePreKey(None) \
                .setTheirBaseKey(parameters.getTheirBaseKey()) \
                .setTheirIdentityKey(parameters.getTheirIdentityKey())
            RatchetingSession.initializeSessionAsBob(sessionState, bobParameters.create())

    @staticmethod
    def initializeSessionAsAlice(sessionState, parameters):
        """
        :type sessionState: SessionState
        :type parameters: AliceAxolotlParameters
        """
        sessionState.setSessionVersion(CiphertextMessage.CURRENT_VERSION)
        sessionState.setRemoteIdentityKey(parameters.getTheirIdentityKey())
        sessionState.setLocalIdentityKey(parameters.getOurIdentityKey().getPublicKey())

        sendingRatchetKey = Curve.generateKeyPair()
        secrets = bytearray()

        secrets.extend(RatchetingSession.getDiscontinuityBytes())

        secrets.extend(Curve.calculateAgreement(parameters.getTheirSignedPreKey(),
                                                parameters.getOurIdentityKey().getPrivateKey()))
        secrets.extend(Curve.calculateAgreement(parameters.getTheirIdentityKey().getPublicKey(),
                                                parameters.getOurBaseKey().getPrivateKey()))
        secrets.extend(Curve.calculateAgreement(parameters.getTheirSignedPreKey(),
                                                parameters.getOurBaseKey().getPrivateKey()))

        if parameters.getTheirOneTimePreKey() is not None:
            secrets.extend(Curve.calculateAgreement(parameters.getTheirOneTimePreKey(),
                                                    parameters.getOurBaseKey().getPrivateKey()))

        derivedKeys = RatchetingSession.calculateDerivedKeys(secrets)
        sendingChain = derivedKeys.getRootKey().createChain(parameters.getTheirRatchetKey(), sendingRatchetKey)

        sessionState.addReceiverChain(parameters.getTheirRatchetKey(), derivedKeys.getChainKey())
        sessionState.setSenderChain(sendingRatchetKey, sendingChain[1])
        sessionState.setRootKey(sendingChain[0])

    @staticmethod
    def initializeSessionAsBob(sessionState, parameters):
        """
        :type sessionState: SessionState
        :type parameters: BobAxolotlParameters
        """
        sessionState.setSessionVersion(CiphertextMessage.CURRENT_VERSION)
        sessionState.setRemoteIdentityKey(parameters.getTheirIdentityKey())
        sessionState.setLocalIdentityKey(parameters.getOurIdentityKey().getPublicKey())

        secrets = bytearray()

        secrets.extend(RatchetingSession.getDiscontinuityBytes())

        secrets.extend(Curve.calculateAgreement(parameters.getTheirIdentityKey().getPublicKey(),
                                                parameters.getOurSignedPreKey().getPrivateKey()))

        secrets.extend(Curve.calculateAgreement(parameters.getTheirBaseKey(),
                                                parameters.getOurIdentityKey().getPrivateKey()))
        secrets.extend(Curve.calculateAgreement(parameters.getTheirBaseKey(),
                                                parameters.getOurSignedPreKey().getPrivateKey()))

        if parameters.getOurOneTimePreKey() is not None:
            secrets.extend(Curve.calculateAgreement(parameters.getTheirBaseKey(),
                                                    parameters.getOurOneTimePreKey().getPrivateKey()))

        derivedKeys = RatchetingSession.calculateDerivedKeys(secrets)
        sessionState.setSenderChain(parameters.getOurRatchetKey(), derivedKeys.getChainKey())
        sessionState.setRootKey(derivedKeys.getRootKey())

    @staticmethod
    def getDiscontinuityBytes():
        return bytearray([0xFF] * 32)

    @staticmethod
    def calculateDerivedKeys(masterSecret):
        kdf = HKDFv3()
        derivedSecretBytes = kdf.deriveSecrets(masterSecret,  bytearray("WhisperText".encode()), 64)
        derivedSecrets = ByteUtil.split(derivedSecretBytes, 32, 32)
        return RatchetingSession.DerivedKeys(RootKey(kdf, derivedSecrets[0]),
                                             ChainKey(kdf, derivedSecrets[1], 0))

    @staticmethod
    def isAlice(ourKey, theirKey):
        """
        :type ourKey: ECPublicKey
        :type theirKey: ECPublicKey
        """
        return ourKey < theirKey

    class DerivedKeys:
        def __init__(self, rootKey, chainKey):
            """
            :type rootKey: RootKey
            :type  chainKey: ChainKey
            """
            self.rootKey = rootKey
            self.chainKey = chainKey

        def getRootKey(self):
            return self.rootKey

        def getChainKey(self):
            return self.chainKey class InvalidKeyException(Exception):
    pass class InvalidKeyIdException(Exception):
    pass class UntrustedIdentityException(Exception):
    def __init__(self, name, identityKey):
        self.name = name
        self.identityKey = identityKey

    def getName(self):
        return self.name

    def getIdentityKey(self):
        return self.identityKey class KeyExchangeMessage:
    INITIATE_FLAG = 0x01
    RESPONSE_FLAG = 0X02
    SIMULTAENOUS_INITIATE_FLAG = 0x04

    def __init__(self, messageVersion=None, sequence=None, flags=None, baseKey=None,
                 baseKeySignature=None, ratchetKey=None, identityKey=None, serialized=None):
        """
        :type messageVersion: int
        :type  sequence: int
        :type flags:int
        :type baseKey: ECPublicKey
        :type baseKeySignature: bytearray
        :type ratchetKey: ECPublicKey
        :type identityKey: IdentityKey
        :type serialized: bytearray
        """
        if serialized:
            try:
                parts = ByteUtil.split(serialized, 1, len(serialized) - 1)
                self.version = ByteUtil.highBitsToInt(parts[0][0])
                self.supportedVersion = ByteUtil.lowBitsToInt(parts[0][0])
                if self.version < CiphertextMessage.CURRENT_VERSION:
                    raise LegacyMessageException("Unsupportmessageed legacy version: %s" % self.version)
                if self.version > CiphertextMessage.CURRENT_VERSION:
                    raise InvalidVersionException("Unkown version: %s" % self.version)
                message = whisperprotos.KeyExchangeMessage()
                message.ParseFromString(bytes(parts[1]))

                if (not message.HasField("id") or not message.HasField("baseKey") or
                        not message.HasField("ratchetKey") or not message.HasField("identityKey") or
                        not message.HasField("baseKeySignature")):
                    raise InvalidMessageException("Some required fields are missing!")

                self.sequence = message.id >> 5
                self.flags = message.id & 0x1f
                self.serialized = serialized
                self.baseKey = Curve.decodePoint(bytearray(message.baseKey), 0)
                self.baseKeySignature = message.baseKeySignature
                self.ratchetKey = Curve.decodePoint(bytearray(message.ratchetKey), 0)
                self.identityKey = IdentityKey(message.identityKey, 0)

            except InvalidKeyException as e:
                raise InvalidMessageException(e)
        else:
            self.supportedVersion = CiphertextMessage.CURRENT_VERSION
            self.version = messageVersion
            self.sequence = sequence
            self.flags = flags
            self.baseKey = baseKey
            self.baseKeySignature = baseKeySignature
            self.ratchetKey = ratchetKey
            self.identityKey = identityKey

            version = [ByteUtil.intsToByteHighAndLow(self.version, self.supportedVersion)]
            keyExchangeMessage = whisperprotos.KeyExchangeMessage()
            keyExchangeMessage.id = (self.sequence << 5) | self.flags
            keyExchangeMessage.baseKey = baseKey.serialize()
            keyExchangeMessage.ratchetKey = ratchetKey.serialize()
            keyExchangeMessage.identityKey = identityKey.serialize()

            if messageVersion >= 3:
                keyExchangeMessage.baseKeySignature = baseKeySignature

            self.serialized = ByteUtil.combine(version, keyExchangeMessage.SerializeToString())

    def getVersion(self):
        return self.version

    def getBaseKey(self):
        return self.baseKey

    def getBaseKeySignature(self):
        return self.baseKeySignature

    def getRatchetKey(self):
        return self.ratchetKey

    def getIdentityKey(self):
        return self.identityKey

    def hasIdentityKey(self):
        return True

    def getMaxVersion(self):
        return self.supportedVersion

    def isResponse(self):
        return ((self.flags & self.__class__.RESPONSE_FLAG) != 0)

    def isInitiate(self):
        return (self.flags & self.__class__.INITIATE_FLAG) != 0

    def isResponseForSimultaneousInitiate(self):
        return (self.flags & self.__class__.SIMULTAENOUS_INITIATE_FLAG) != 0

    def getFlags(self):
        return self.flags

    def getSequence(self):
        return self.sequence

    def serialize(self):
        return self.serialized class CiphertextMessage(object):
    __metaclass__ = abc.ABCMeta

    UNSUPPORTED_VERSION = 1
    CURRENT_VERSION = 3

    WHISPER_TYPE = 2
    PREKEY_TYPE = 3
    SENDERKEY_TYPE = 4
    SENDERKEY_DISTRIBUTION_TYPE = 5

    # This should be the worst case (worse than V2). So not always accurate, but good enough for padding.
    ENCRYPTED_MESSAGE_OVERHEAD = 53

    @abc.abstractmethod
    def serialize(self):
        return

    @abc.abstractmethod
    def getType(self):
        return class StaleKeyExchangeException(Exception):
    pass class Medium:
    MAX_VALUE = 0xFFFFFF class KeyHelper:
    def __init__(self):
        pass

    @staticmethod
    def generateIdentityKeyPair():
        """
        Generate an identity key pair.  Clients should only do this once,
        at install time.
        @return the generated IdentityKeyPair.
        """
        keyPair = Curve.generateKeyPair()
        publicKey = IdentityKey(keyPair.getPublicKey())
        serialized = '0a21056e8936e8367f768a7bba008ade7cf58407bdc7a6aae293e2c' \
                     'b7c06668dcd7d5e12205011524f0c15467100dd603e0d6020f4d293' \
                     'edfbcd82129b14a88791ac81365c'
        serialized = binascii.unhexlify(serialized.encode())
        identityKeyPair = IdentityKeyPair(publicKey, keyPair.getPrivateKey())
        return identityKeyPair
        # return IdentityKeyPair(serialized=serialized)

    @staticmethod
    def generateRegistrationId(extended_range=False):
        """
        Generate a registration ID.  Clients should only do this once,
        at install time.
        :param extended_range: By default (false), the generated registration ID is sized to require the minimal
        possible protobuf encoding overhead. Specify true if the caller needs the full range of MAX_INT at the cost
        of slightly higher encoding overhead.
        """
        if extended_range:
            regId = KeyHelper.getRandomSequence(2147483646) + 1
        else:
            regId = KeyHelper.getRandomSequence(16380) + 1

        return regId

    @staticmethod
    def getRandomSequence(max):
        return SystemRandom().randrange(max)

    @staticmethod
    def generatePreKeys(start, count):
        """
        Generate a list of PreKeys.  Clients should do this at install time, and
        subsequently any time the list of PreKeys stored on the server runs low.

        PreKey IDs are shorts, so they will eventually be repeated.  Clients should
        store PreKeys in a circular buffer, so that they are repeated as infrequently
        as possible.

        @param start The starting PreKey ID, inclusive.
        @param count The number of PreKeys to generate.
        @return the list of generated PreKeyRecords.
        """
        results = []
        start -= 1
        for i in range(0, count):
            preKeyId = ((start + i) % (Medium.MAX_VALUE - 1)) + 1
            results.append(PreKeyRecord(preKeyId, Curve.generateKeyPair()))

        return results

    @staticmethod
    def generateSignedPreKey(identityKeyPair, signedPreKeyId):
        keyPair = Curve.generateKeyPair()
        signature = Curve.calculateSignature(identityKeyPair.getPrivateKey(), keyPair.getPublicKey().serialize())

        spk = SignedPreKeyRecord(signedPreKeyId, int(round(time.time() * 1000)), keyPair, signature)

        return spk

    @staticmethod
    def generateSenderSigningKey():
        return Curve.generateKeyPair()

    @staticmethod
    def generateSenderKey():
        return os.urandom(32)

    @staticmethod
    def generateSenderKeyId():
        return KeyHelper.getRandomSequence(2147483647)
import logging
from .ecc.curve import Curve
from .ratchet.aliceaxolotlparameters import AliceAxolotlParameters
from .ratchet.bobaxolotlparamaters import BobAxolotlParameters
from .ratchet.symmetricaxolotlparameters import SymmetricAxolotlParameters
from .ratchet.ratchetingsession import RatchetingSession
from .invalidkeyexception import InvalidKeyException
from .invalidkeyidexception import InvalidKeyIdException
from .untrustedidentityexception import UntrustedIdentityException
from .protocol.keyexchangemessage import KeyExchangeMessage
from .protocol.ciphertextmessage import CiphertextMessage
from .statekeyexchangeexception import StaleKeyExchangeException
from .util.medium import Medium
from .util.keyhelper import KeyHelper

        if message.getPreKeyId() is not None and message.getPreKeyId() != Medium.MAX_VALUE:
            return message.getPreKeyId()
        else:
            return None

    def processPreKeyBundle(self, preKey):
        """
        :type preKey: PreKeyBundle
        """
        if not self.identityKeyStore.isTrustedIdentity(self.recipientId, preKey.getIdentityKey()):
            raise UntrustedIdentityException(self.recipientId, preKey.getIdentityKey())

        if preKey.getSignedPreKey() is not None and\
            not Curve.verifySignature(preKey.getIdentityKey().getPublicKey(),
                                      preKey.getSignedPreKey().serialize(),
                                      preKey.getSignedPreKeySignature()):
            raise InvalidKeyException("Invalid signature on device key!")

        if preKey.getSignedPreKey() is None:
            raise InvalidKeyException("No signed prekey!!")

        sessionRecord = self.sessionStore.loadSession(self.recipientId, self.deviceId)
        ourBaseKey = Curve.generateKeyPair()
        theirSignedPreKey = preKey.getSignedPreKey()
        theirOneTimePreKey = preKey.getPreKey()
        theirOneTimePreKeyId = preKey.getPreKeyId() if theirOneTimePreKey is not None else None

        parameters = AliceAxolotlParameters.newBuilder()

        parameters.setOurBaseKey(ourBaseKey)\
            .setOurIdentityKey(self.identityKeyStore.getIdentityKeyPair())\
            .setTheirIdentityKey(preKey.getIdentityKey())\
            .setTheirSignedPreKey(theirSignedPreKey)\
            .setTheirRatchetKey(theirSignedPreKey)\
            .setTheirOneTimePreKey(theirOneTimePreKey)

        if not sessionRecord.isFresh():
            sessionRecord.archiveCurrentState()

        RatchetingSession.initializeSessionAsAlice(sessionRecord.getSessionState(), parameters.create())

        sessionRecord.getSessionState().setUnacknowledgedPreKeyMessage(theirOneTimePreKeyId,
                                                                       preKey.getSignedPreKeyId(),
                                                                       ourBaseKey.getPublicKey())
        sessionRecord.getSessionState().setLocalRegistrationId(self.identityKeyStore.getLocalRegistrationId())
        sessionRecord.getSessionState().setRemoteRegistrationId(preKey.getRegistrationId())
        sessionRecord.getSessionState().setAliceBaseKey(ourBaseKey.getPublicKey().serialize())
        self.sessionStore.storeSession(self.recipientId, self.deviceId, sessionRecord)
        self.identityKeyStore.saveIdentity(self.recipientId, preKey.getIdentityKey())

    def processKeyExchangeMessage(self, keyExchangeMessage):

        if not self.identityKeyStore.isTrustedIdentity(self.recipientId, keyExchangeMessage.getIdentityKey()):
            raise UntrustedIdentityException(self.recipientId, keyExchangeMessage.getIdentityKey())

        responseMessage = None

        if keyExchangeMessage.isInitiate():
            responseMessage = self.processInitiate(keyExchangeMessage)
        else:
            self.processResponse(keyExchangeMessage)

        return responseMessage

    def processInitiate(self, keyExchangeMessage):
        flags = KeyExchangeMessage.RESPONSE_FLAG
        sessionRecord = self.sessionStore.loadSession(self.recipientId, self.deviceId)

        if not Curve.verifySignature(
                keyExchangeMessage.getIdentityKey().getPublicKey(),
                keyExchangeMessage.getBaseKey().serialize(),
                keyExchangeMessage.getBaseKeySignature()):
            raise InvalidKeyException("Bad signature!")

        builder = SymmetricAxolotlParameters.newBuilder()
        if not sessionRecord.getSessionState().hasPendingKeyExchange():
            builder.setOurIdentityKey(self.identityKeyStore.getIdentityKeyPair())\
                .setOurBaseKey(Curve.generateKeyPair())\
                .setOurRatchetKey(Curve.generateKeyPair())
        else:
            builder.setOurIdentityKey(sessionRecord.getSessionState().getPendingKeyExchangeIdentityKey())\
                .setOurBaseKey(sessionRecord.getSessionState().getPendingKeyExchangeBaseKey())\
                .setOurRatchetKey(sessionRecord.getSessionState().getPendingKeyExchangeRatchetKey())
            flags |= KeyExchangeMessage.SIMULTAENOUS_INITIATE_FLAG

        builder.setTheirBaseKey(keyExchangeMessage.getBaseKey())\
            .setTheirRatchetKey(keyExchangeMessage.getRatchetKey())\
            .setTheirIdentityKey(keyExchangeMessage.getIdentityKey())

        parameters = builder.create()

        if not sessionRecord.isFresh():
            sessionRecord.archiveCurrentState()

        RatchetingSession.initializeSession(sessionRecord.getSessionState(), parameters)

        self.sessionStore.storeSession(self.recipientId, self.deviceId, sessionRecord)
        self.identityKeyStore.saveIdentity(self.recipientId, keyExchangeMessage.getIdentityKey())

        baseKeySignature = Curve.calculateSignature(parameters.getOurIdentityKey().getPrivateKey(),
                                                    parameters.getOurBaseKey().getPublicKey().serialize())

        return KeyExchangeMessage(sessionRecord.getSessionState().getSessionVersion(),
                                  keyExchangeMessage.getSequence(), flags,
                                  parameters.getOurBaseKey().getPublicKey(),
                                  baseKeySignature, parameters.getOurRatchetKey().getPublicKey(),
                                  parameters.getOurIdentityKey().getPublicKey())

    def processResponse(self, keyExchangeMessage):
        sessionRecord = self.sessionStore.loadSession(self.recipientId, self.deviceId)
        sessionState = sessionRecord.getSessionState()
        hasPendingKeyExchange = sessionState.hasPendingKeyExchange()
        isSimultaneousInitiateResponse = keyExchangeMessage.isResponseForSimultaneousInitiate()

        if not hasPendingKeyExchange \
                or sessionState.getPendingKeyExchangeSequence() != keyExchangeMessage.getSequence():
            logger.warn("No matching sequence for response. "
                        "Is simultaneous initiate response: %s" % isSimultaneousInitiateResponse)
            if not isSimultaneousInitiateResponse:

                raise StaleKeyExchangeException()