diff --git a/nucypher/core.py b/nucypher/core.py index ee0db10be..9bb15ca4d 100644 --- a/nucypher/core.py +++ b/nucypher/core.py @@ -15,7 +15,7 @@ You should have received a copy of the GNU Affero General Public License along with nucypher. If not, see . """ -from typing import Optional, Sequence, Callable, Dict, Tuple, List +from typing import Optional, Sequence, Callable, Dict, Tuple, List, Iterable from bytestring_splitter import ( BytestringSplitter, @@ -34,6 +34,7 @@ from nucypher.crypto.splitters import ( key_splitter, kfrag_splitter, cfrag_splitter, + checksum_address_splitter, ) from nucypher.crypto.signing import InvalidSignature import nucypher.crypto.umbral_adapter as umbral # need it to mock `umbral.encrypt` @@ -663,3 +664,45 @@ class ReencryptionResponse(Versioned): cfrags = cfrag_splitter.repeat(cfrags_bytes) return cls(cfrags, signature) + + +class RetrievalKit(Versioned): + """ + An object encapsulating the information necessary for retrieval of cfrags from Ursulas. + Contains the capsule and the checksum addresses of Ursulas from which the requester + already received cfrags. + """ + + @classmethod + def from_message_kit(cls, message_kit: MessageKit) -> 'RetrievalKit': + return cls(message_kit.capsule, set()) + + def __init__(self, capsule: Capsule, queried_addresses: Iterable[ChecksumAddress]): + self.capsule = capsule + # Can store cfrags too, if we're worried about Ursulas supplying duplicate ones. + self.queried_addresses = set(queried_addresses) + + def _payload(self) -> bytes: + return (bytes(self.capsule) + + b''.join(to_canonical_address(address) for address in self.queried_addresses)) + + @classmethod + def _brand(cls) -> bytes: + return b'RKit' + + @classmethod + def _version(cls) -> Tuple[int, int]: + return 1, 0 + + @classmethod + def _old_version_handlers(cls) -> Dict: + return {} + + @classmethod + def _from_bytes_current(cls, data): + capsule, remainder = capsule_splitter(data, return_remainder=True) + if remainder: + addresses_as_bytes = checksum_address_splitter.repeat(remainder) + else: + addresses_as_bytes = () + return cls(capsule, set(to_checksum_address(address) for address in addresses_as_bytes)) diff --git a/nucypher/network/retrieval.py b/nucypher/network/retrieval.py index 225451a91..97ce4886a 100644 --- a/nucypher/network/retrieval.py +++ b/nucypher/network/retrieval.py @@ -26,6 +26,7 @@ from nucypher.core import ( TreasureMap, ReencryptionResponse, ReencryptionRequest, + RetrievalKit, ) from nucypher.crypto.signing import InvalidSignature @@ -37,8 +38,7 @@ from nucypher.crypto.umbral_adapter import ( ) from nucypher.network.exceptions import NodeSeemsToBeDown from nucypher.network.nodes import Learner -from nucypher.policy.kits import RetrievalKit, RetrievalResult -from nucypher.utilities.versioning import Versioned +from nucypher.policy.kits import RetrievalResult class RetrievalPlan: diff --git a/nucypher/policy/kits.py b/nucypher/policy/kits.py index f80db7391..773596a1e 100644 --- a/nucypher/policy/kits.py +++ b/nucypher/policy/kits.py @@ -16,63 +16,16 @@ along with nucypher. If not, see . """ -from typing import Dict, Iterable, Set, Union, Tuple +from typing import Dict, Set, Union from eth_typing import ChecksumAddress -from eth_utils import to_checksum_address, to_canonical_address -from nucypher.core import MessageKit +from nucypher.core import MessageKit, RetrievalKit -from nucypher.crypto.splitters import ( - capsule_splitter, - checksum_address_splitter, -) from nucypher.crypto.umbral_adapter import PublicKey, VerifiedCapsuleFrag, Capsule, SecretKey from nucypher.utilities.versioning import Versioned -class RetrievalKit(Versioned): - """ - An object encapsulating the information necessary for retrieval of cfrags from Ursulas. - Contains the capsule and the checksum addresses of Ursulas from which the requester - already received cfrags. - """ - - @classmethod - def from_message_kit(cls, message_kit: MessageKit) -> 'RetrievalKit': - return cls(message_kit.capsule, set()) - - def __init__(self, capsule: Capsule, queried_addresses: Iterable[ChecksumAddress]): - self.capsule = capsule - # Can store cfrags too, if we're worried about Ursulas supplying duplicate ones. - self.queried_addresses = set(queried_addresses) - - def _payload(self) -> bytes: - return (bytes(self.capsule) + - b''.join(to_canonical_address(address) for address in self.queried_addresses)) - - @classmethod - def _brand(cls) -> bytes: - return b'RKit' - - @classmethod - def _version(cls) -> Tuple[int, int]: - return 1, 0 - - @classmethod - def _old_version_handlers(cls) -> Dict: - return {} - - @classmethod - def _from_bytes_current(cls, data): - capsule, remainder = capsule_splitter(data, return_remainder=True) - if remainder: - addresses_as_bytes = checksum_address_splitter.repeat(remainder) - else: - addresses_as_bytes = () - return cls(capsule, set(to_checksum_address(address) for address in addresses_as_bytes)) - - class PolicyMessageKit: @classmethod diff --git a/nucypher/utilities/porter/control/interfaces.py b/nucypher/utilities/porter/control/interfaces.py index 33a6cb479..7ea53565d 100644 --- a/nucypher/utilities/porter/control/interfaces.py +++ b/nucypher/utilities/porter/control/interfaces.py @@ -18,11 +18,10 @@ from typing import List, Optional from eth_typing import ChecksumAddress -from nucypher.core import TreasureMap +from nucypher.core import TreasureMap, RetrievalKit from nucypher.control.interfaces import ControlInterface, attach_schema from nucypher.crypto.umbral_adapter import PublicKey -from nucypher.policy.kits import RetrievalKit from nucypher.utilities.porter.control.specifications import porter_schema diff --git a/nucypher/utilities/porter/control/specifications/fields/retrieve.py b/nucypher/utilities/porter/control/specifications/fields/retrieve.py index 05fa0e71f..5df040036 100644 --- a/nucypher/utilities/porter/control/specifications/fields/retrieve.py +++ b/nucypher/utilities/porter/control/specifications/fields/retrieve.py @@ -16,11 +16,12 @@ """ from marshmallow import fields +from nucypher.core import RetrievalKit as RetrievalKitClass + from nucypher.control.specifications.base import BaseSchema from nucypher.control.specifications.exceptions import InvalidInputData from nucypher.control.specifications.fields import Base64BytesRepresentation from nucypher.crypto.umbral_adapter import CapsuleFrag as CapsuleFragClass -from nucypher.policy.kits import RetrievalKit as RetrievalKitClass from nucypher.utilities.porter.control.specifications.fields import UrsulaChecksumAddress diff --git a/nucypher/utilities/porter/porter.py b/nucypher/utilities/porter/porter.py index 9517f6391..d5fb3b05f 100644 --- a/nucypher/utilities/porter/porter.py +++ b/nucypher/utilities/porter/porter.py @@ -21,7 +21,7 @@ from constant_sorrow.constants import NO_BLOCKCHAIN_CONNECTION, NO_CONTROL_PROTO from eth_typing import ChecksumAddress from flask import request, Response -from nucypher.core import TreasureMap +from nucypher.core import TreasureMap, RetrievalKit from nucypher.blockchain.eth.agents import ContractAgency, StakingEscrowAgent from nucypher.blockchain.eth.interfaces import BlockchainInterfaceFactory @@ -32,7 +32,7 @@ from nucypher.crypto.powers import DecryptingPower from nucypher.crypto.umbral_adapter import PublicKey from nucypher.network.nodes import Learner from nucypher.network.retrieval import RetrievalClient -from nucypher.policy.kits import RetrievalKit, RetrievalResult +from nucypher.policy.kits import RetrievalResult from nucypher.policy.reservoir import ( make_federated_staker_reservoir, make_decentralized_staker_reservoir, diff --git a/tests/acceptance/porter/control/test_porter_web_control_blockchain.py b/tests/acceptance/porter/control/test_porter_web_control_blockchain.py index 58534aca3..68706cae6 100644 --- a/tests/acceptance/porter/control/test_porter_web_control_blockchain.py +++ b/tests/acceptance/porter/control/test_porter_web_control_blockchain.py @@ -20,9 +20,11 @@ import os from base64 import b64encode from urllib.parse import urlencode +from nucypher.core import RetrievalKit + from nucypher.characters.lawful import Enrico from nucypher.crypto.powers import DecryptingPower -from nucypher.policy.kits import PolicyMessageKit, RetrievalResult, RetrievalKit +from nucypher.policy.kits import PolicyMessageKit, RetrievalResult from nucypher.utilities.porter.control.specifications.fields import RetrievalResultSchema, RetrievalKit as RetrievalKitField from tests.utils.middleware import MockRestMiddleware from tests.utils.policy import retrieval_request_setup, retrieval_params_decode_from_rest diff --git a/tests/integration/characters/test_bob_handles_frags.py b/tests/integration/characters/test_bob_handles_frags.py index 181e86684..909d9ce7b 100644 --- a/tests/integration/characters/test_bob_handles_frags.py +++ b/tests/integration/characters/test_bob_handles_frags.py @@ -19,10 +19,11 @@ import pytest import pytest_twisted from twisted.internet import threads +from nucypher.core import RetrievalKit + from nucypher.characters.lawful import Enrico, Bob from nucypher.config.constants import TEMPORARY_DOMAIN from nucypher.network.retrieval import RetrievalClient -from nucypher.policy.kits import RetrievalKit from tests.utils.middleware import MockRestMiddleware, NodeIsDownMiddleware diff --git a/tests/integration/porter/control/test_porter_web_control_federated.py b/tests/integration/porter/control/test_porter_web_control_federated.py index 3d9aef635..0dd973ddf 100644 --- a/tests/integration/porter/control/test_porter_web_control_federated.py +++ b/tests/integration/porter/control/test_porter_web_control_federated.py @@ -20,9 +20,11 @@ import json from base64 import b64encode from urllib.parse import urlencode +from nucypher.core import RetrievalKit + from nucypher.characters.lawful import Enrico from nucypher.crypto.powers import DecryptingPower -from nucypher.policy.kits import PolicyMessageKit, RetrievalResult, RetrievalKit +from nucypher.policy.kits import PolicyMessageKit, RetrievalResult from nucypher.utilities.porter.control.specifications.fields import RetrievalResultSchema, RetrievalKit as RetrievalKitField from tests.utils.policy import retrieval_request_setup, retrieval_params_decode_from_rest diff --git a/tests/unit/test_porter.py b/tests/unit/test_porter.py index 71e655d63..77ed92277 100644 --- a/tests/unit/test_porter.py +++ b/tests/unit/test_porter.py @@ -20,10 +20,11 @@ from base64 import b64encode import pytest +from nucypher.core import RetrievalKit as RetrievalKitClass + from nucypher.control.specifications.exceptions import InvalidInputData from nucypher.control.specifications.fields import StringList from nucypher.crypto.umbral_adapter import SecretKey, encrypt -from nucypher.policy.kits import RetrievalKit as RetrievalKitClass from nucypher.utilities.porter.control.specifications.fields import HRAC, UrsulaChecksumAddress from nucypher.utilities.porter.control.specifications.fields.retrieve import RetrievalKit diff --git a/tests/utils/policy.py b/tests/utils/policy.py index a39f9b154..5768d45fc 100644 --- a/tests/utils/policy.py +++ b/tests/utils/policy.py @@ -20,12 +20,11 @@ import random import string from typing import Dict, Tuple -from nucypher.core import MessageKit +from nucypher.core import MessageKit, RetrievalKit from nucypher.characters.control.specifications.fields import Key, TreasureMap from nucypher.characters.lawful import Enrico from nucypher.crypto.powers import DecryptingPower -from nucypher.policy.kits import RetrievalKit from nucypher.utilities.porter.control.specifications.fields import RetrievalKit as RetrievalKitField