Move RetrievalKit to core.py

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 <https://www.gnu.org/licenses/>.
 """
 
-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))

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:

nucypher/policy/kits.py

@@ -16,63 +16,16 @@ along with nucypher.  If not, see <https://www.gnu.org/licenses/>.
 """
 
 
-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

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
 
 

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
 
 

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,

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

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
 

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
 

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
 

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