From 9170b7013547b6e7259f8270b4477649a8f8e6aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Thu, 31 May 2018 14:27:20 +0200
Subject: [PATCH 01/11] Add __eq__ method to UmbralParameters

---
 umbral/params.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/umbral/params.py b/umbral/params.py
index 3fac503a..e6ec7f4f 100644
--- a/umbral/params.py
+++ b/umbral/params.py
@@ -21,3 +21,14 @@ def __init__(self, curve: ec.EllipticCurve):
 
         parameters_seed = b'NuCypherKMS/UmbralParameters/'
         self.u = unsafe_hash_to_point(g_bytes, self, parameters_seed + b'u')
+
+    def __eq__(self, other):
+
+        self_curve_nid = backend._elliptic_curve_to_nid(self.curve)
+        other_curve_nid = backend._elliptic_curve_to_nid(other.curve)
+
+        # TODO: This is not comparing the order, which currently is an OpenSSL pointer
+        self_attributes = self_curve_nid, self.g, self.CURVE_KEY_SIZE_BYTES, self.u
+        others_attributes = other_curve_nid, other.g, other.CURVE_KEY_SIZE_BYTES, other.u
+
+        return self_attributes == others_attributes
\ No newline at end of file

From 03b4138229cbb5c8b75822975b4adc36c0c48687 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Thu, 31 May 2018 14:31:28 +0200
Subject: [PATCH 02/11] Add params when deserializing Capsule

---
 umbral/pre.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/umbral/pre.py b/umbral/pre.py
index a49822f7..9fef7087 100644
--- a/umbral/pre.py
+++ b/umbral/pre.py
@@ -91,15 +91,18 @@ class NotValid(ValueError):
         """
 
     @classmethod
-    def from_bytes(cls, capsule_bytes: bytes, curve: ec.EllipticCurve = None):
+    def from_bytes(cls, capsule_bytes: bytes, params: UmbralParameters = None):
         """
         Instantiates a Capsule object from the serialized data.
         """
-        curve = curve if curve is not None else default_curve()
+        params = params if params is not None else default_params()
+        curve = params.curve
+
         bn_size = CurveBN.expected_bytes_length(curve)
         point_size = Point.expected_bytes_length(curve)
 
-        if len(capsule_bytes) == cls.expected_bytes_length(curve, activated=True):
+        capsule_bytes_length = len(capsule_bytes)
+        if capsule_bytes_length == cls.expected_bytes_length(curve, activated=True):
             splitter = BytestringSplitter(
                 (Point, point_size),  # point_e
                 (Point, point_size),  # point_v
@@ -108,15 +111,17 @@ def from_bytes(cls, capsule_bytes: bytes, curve: ec.EllipticCurve = None):
                 (Point, point_size),  # point_v_prime
                 (Point, point_size)  # point_noninteractive
             )
-        else:
+        elif capsule_bytes_length == cls.expected_bytes_length(curve, activated=False):
             splitter = BytestringSplitter(
                 (Point, point_size),  # point_e
                 (Point, point_size),  # point_v
                 (CurveBN, bn_size)  # bn_sig
             )
+        else:
+            raise ValueError("Byte string does not have a valid length for a Capsule")
 
         components = splitter(capsule_bytes)
-        return cls(*components)
+        return cls(*components, params=params)
 
     def _set_cfrag_correctness_key(self, key_type, key: UmbralPublicKey):
         if key_type not in ("delegating", "receiving", "verifying"): 

From c98f723eec301764d811a27b56832b663c3c3d4d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Thu, 31 May 2018 14:32:04 +0200
Subject: [PATCH 03/11] Check that params match when setting correctness keys

---
 umbral/pre.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/umbral/pre.py b/umbral/pre.py
index 9fef7087..5482849e 100644
--- a/umbral/pre.py
+++ b/umbral/pre.py
@@ -57,8 +57,8 @@ def __init__(self,
                 "Passing both is also fine.")
 
         self._cfrag_correctness_keys = {"delegating": delegating_pubkey,
-                                     "receiving": receiving_pubkey,
-                                     "verifying": verifying_pubkey}
+                                        "receiving": receiving_pubkey,
+                                        "verifying": verifying_pubkey}
 
         self._point_e = point_e
         self._point_v = point_v
@@ -132,6 +132,8 @@ def _set_cfrag_correctness_key(self, key_type, key: UmbralPublicKey):
         if current_key is None:
             if key is None:
                 raise TypeError("The {} key is not set and you didn't pass one.".format(key_type))
+            elif self._umbral_params != key.params:
+                raise TypeError("You are trying to set a key with different UmbralParameters.")
             else:
                 self._cfrag_correctness_keys[key_type] = key
                 return True

From 09bed343059f585f506e0dde527c901cd988166a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Thu, 31 May 2018 14:33:11 +0200
Subject: [PATCH 04/11] Remove keypairs fixtures from test_simple_api

---
 tests/test_simple_api.py | 23 ++++++++---------------
 1 file changed, 8 insertions(+), 15 deletions(-)

diff --git a/tests/test_simple_api.py b/tests/test_simple_api.py
index 39a058c1..605a8369 100644
--- a/tests/test_simple_api.py
+++ b/tests/test_simple_api.py
@@ -15,18 +15,21 @@
 
 
 @pytest.mark.parametrize("N, M", parameters)
-def test_simple_api(alices_keys, bobs_keys, N, M, curve=default_curve()):
+def test_simple_api(N, M, curve=default_curve()):
     """Manually injects umbralparameters for multi-curve testing."""
 
     params = UmbralParameters(curve=curve)
 
-    delegating_privkey, signing_privkey = alices_keys
+    delegating_privkey = keys.UmbralPrivateKey.gen_key(params=params)
     delegating_pubkey = delegating_privkey.get_pubkey()
-    signing_pubkey = signing_privkey.get_pubkey()
 
-    receiving_privkey, receiving_pubkey = bobs_keys
+    signing_privkey = keys.UmbralPrivateKey.gen_key(params=params)
+    signing_pubkey = signing_privkey.get_pubkey()
     signer = Signer(signing_privkey)
 
+    receiving_privkey = keys.UmbralPrivateKey.gen_key(params=params)
+    receiving_pubkey = receiving_privkey.get_pubkey()
+
     plain_data = b'peace at dawn'
     ciphertext, capsule = pre.encrypt(delegating_pubkey, plain_data, params=params)
 
@@ -52,17 +55,7 @@ def test_simple_api(alices_keys, bobs_keys, N, M, curve=default_curve()):
 @pytest.mark.parametrize("curve", secp_curves)
 @pytest.mark.parametrize("N, M", parameters)
 def test_simple_api_on_multiple_curves(N, M, curve):
-    params = UmbralParameters(curve=curve)
-
-    delegating_privkey = keys.UmbralPrivateKey.gen_key(params=params)
-    signing_privkey = keys.UmbralPrivateKey.gen_key(params=params)
-    alices_keys = delegating_privkey, signing_privkey
-
-    receiving_privkey = keys.UmbralPrivateKey.gen_key(params=params)
-    receiving_pubkey = receiving_privkey.get_pubkey()
-    bobs_keys = receiving_privkey, receiving_pubkey
-    
-    test_simple_api(alices_keys, bobs_keys, N, M, curve)
+    test_simple_api(N, M, curve)
 
 
 def test_public_key_encryption(alices_keys):

From c736370bcc3b16426132f02996d5ff6ef505d0be Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Fri, 1 Jun 2018 12:28:45 +0200
Subject: [PATCH 05/11] Remove params as an argument when possible

* Instead, takes params from Capsules or UmbralPublicKeys
* Makes params required in some places (Capsule.init, Capsule.from_bytes, etc)
* Removes pre.CHACHA20_KEY_SIZE constant and use dem.DEM_KEYSIZE instead
---
 umbral/_pre.py      |  13 ++---
 umbral/curvebn.py   |   4 +-
 umbral/fragments.py |  16 +++----
 umbral/keys.py      |  38 +++++++--------
 umbral/point.py     |   3 +-
 umbral/pre.py       | 113 ++++++++++++++++++++------------------------
 6 files changed, 84 insertions(+), 103 deletions(-)

diff --git a/umbral/_pre.py b/umbral/_pre.py
index 72edcb66..cf3ac9a0 100644
--- a/umbral/_pre.py
+++ b/umbral/_pre.py
@@ -7,10 +7,9 @@
 def prove_cfrag_correctness(cfrag: "CapsuleFrag",
                             kfrag: "KFrag",
                             capsule: "Capsule",
-                            metadata: bytes = None,
-                            params: UmbralParameters = None
+                            metadata: bytes = None
                             ) -> "CorrectnessProof":
-    params = params if params is not None else default_params()
+    params = capsule._umbral_params
 
     rk = kfrag._bn_key
     t = CurveBN.gen_rand(params.curve)
@@ -50,9 +49,8 @@ def assess_cfrag_correctness(cfrag,
                              capsule: "Capsule",
                              delegating_point,
                              signing_pubkey,
-                             receiving_point,
-                             params: UmbralParameters = None):
-    params = params if params is not None else default_params()
+                             receiving_point):
+    params = capsule._umbral_params
 
     ####
     ## Here are the formulaic constituents shared with `prove_cfrag_correctness`.
@@ -107,9 +105,8 @@ def verify_kfrag(kfrag,
                  delegating_point,
                  signing_pubkey,
                  receiving_point,
-                 params: UmbralParameters = None
+                 params: UmbralParameters
                  ):
-    params = params if params is not None else default_params()
 
     u = params.u
 
diff --git a/umbral/curvebn.py b/umbral/curvebn.py
index 657da3f9..2c71ce8f 100644
--- a/umbral/curvebn.py
+++ b/umbral/curvebn.py
@@ -7,6 +7,7 @@
 from umbral import openssl
 from umbral.config import default_curve, default_params
 from umbral.utils import get_curve_keysize_bytes
+from umbral.params import UmbralParameters
 
 
 class CurveBN(object):
@@ -84,8 +85,7 @@ def from_int(cls, num, curve: ec.EllipticCurve=None):
         return cls(conv_bn, curve_nid, group, order)
 
     @classmethod
-    def hash(cls, *crypto_items, params=None):
-        params = params if params is not None else default_params()
+    def hash(cls, *crypto_items, params: UmbralParameters):
 
         curve_nid = backend._elliptic_curve_to_nid(params.curve)
         order = openssl._get_ec_order_by_curve_nid(curve_nid)
diff --git a/umbral/fragments.py b/umbral/fragments.py
index fb48f9dc..b9d8f259 100644
--- a/umbral/fragments.py
+++ b/umbral/fragments.py
@@ -68,10 +68,11 @@ def to_bytes(self):
         return self._id + key + ni + commitment + xcoord + signature
 
     def verify(self,
-               signing_pubkey: UmbralPublicKey,
+               signing_pubkey,
                delegating_pubkey: UmbralPublicKey,
-               receiving_pubkey: UmbralPublicKey,
-               params: UmbralParameters = None):
+               receiving_pubkey: UmbralPublicKey):
+
+        params = delegating_pubkey.params
 
         return verify_kfrag(self,
                             delegating_pubkey.point_key,
@@ -148,7 +149,7 @@ def to_bytes(self) -> bytes:
 
         return result
 
-    def _bn_keytes__(self):
+    def _bytes__(self):
         return self.to_bytes()
 
 
@@ -222,9 +223,8 @@ def to_bytes(self):
     def verify_correctness(self,
                            capsule: "Capsule",
                            delegating_pubkey: UmbralPublicKey,
-                           signing_pubkey: UmbralPublicKey,
-                           receiving_pubkey: UmbralPublicKey,
-                           params: UmbralParameters = None):
+                           signing_pubkey,
+                           receiving_pubkey: UmbralPublicKey):
         if not all((delegating_pubkey, signing_pubkey, receiving_pubkey)):
             raise TypeError("Need all three keys to verify correctness.")
 
@@ -232,7 +232,7 @@ def verify_correctness(self,
         pubkey_b_point = receiving_pubkey.point_key
 
         return assess_cfrag_correctness(self, capsule, pubkey_a_point,
-                                        signing_pubkey, pubkey_b_point, params)
+                                        signing_pubkey, pubkey_b_point)
 
     def attach_proof(self, e2, v2, u1, u2, z3, kfrag_signature, metadata):
         self.proof = CorrectnessProof(point_e2=e2,
diff --git a/umbral/keys.py b/umbral/keys.py
index 3ad2224f..59545151 100644
--- a/umbral/keys.py
+++ b/umbral/keys.py
@@ -20,19 +20,16 @@
 
     
 class UmbralPrivateKey(object):
-    def __init__(self, bn_key: CurveBN, params: UmbralParameters=None):
+    def __init__(self, bn_key: CurveBN, params: UmbralParameters):
         """
         Initializes an Umbral private key.
         """
-        if params is None:
-            params = default_params()
-
         self.params = params
         self.bn_key = bn_key
-        self.pubkey = UmbralPublicKey(self.bn_key * self.params.g, params=params)
+        self.pubkey = UmbralPublicKey(self.bn_key * params.g, params=params)
 
     @classmethod
-    def gen_key(cls, params: UmbralParameters=None):
+    def gen_key(cls, params: UmbralParameters = None):
         """
         Generates a private key and returns it.
         """
@@ -43,9 +40,9 @@ def gen_key(cls, params: UmbralParameters=None):
         return cls(bn_key, params)
 
     @classmethod
-    def from_bytes(cls, key_bytes: bytes, params: UmbralParameters=None,
-                   password: bytes=None, _scrypt_cost: int=20,
-                   decoder: Callable=None):
+    def from_bytes(cls, key_bytes: bytes, params: UmbralParameters = None,
+                   password: bytes = None, _scrypt_cost: int = 20,
+                   decoder: Callable = None):
         """
         Loads an Umbral private key from bytes.
         Optionally, allows a decoder function to be passed as a param to decode
@@ -82,8 +79,8 @@ def from_bytes(cls, key_bytes: bytes, params: UmbralParameters=None,
         bn_key = CurveBN.from_bytes(key_bytes, params.curve)
         return cls(bn_key, params)
 
-    def to_bytes(self, password: bytes=None, _scrypt_cost: int=20,
-                 encoder: Callable=None):
+    def to_bytes(self, password: bytes = None, _scrypt_cost: int = 20,
+                 encoder: Callable = None):
         """
         Returns an Umbral private key as bytes optional symmetric encryption
         via nacl's Salsa20-Poly1305 and Scrypt key derivation. If a password
@@ -163,13 +160,10 @@ def to_cryptography_privkey(self):
 
 
 class UmbralPublicKey(object):
-    def __init__(self, point_key, params: UmbralParameters=None):
+    def __init__(self, point_key, params: UmbralParameters):
         """
         Initializes an Umbral public key.
         """
-        if params is None:
-            params = default_params()
-
         self.params = params
 
         if not isinstance(point_key, Point):
@@ -178,8 +172,8 @@ def __init__(self, point_key, params: UmbralParameters=None):
         self.point_key = point_key
 
     @classmethod
-    def from_bytes(cls, key_bytes: bytes, params: UmbralParameters=None,
-                   decoder: Callable=None):
+    def from_bytes(cls, key_bytes: bytes, params: UmbralParameters = None,
+                   decoder: Callable = None):
         """
         Loads an Umbral public key from bytes.
         Optionally, if an decoder function is provided it will be used to decode
@@ -194,7 +188,7 @@ def from_bytes(cls, key_bytes: bytes, params: UmbralParameters=None,
         point_key = Point.from_bytes(key_bytes, params.curve)
         return cls(point_key, params)
 
-    def to_bytes(self, encoder: Callable=None):
+    def to_bytes(self, encoder: Callable = None):
         """
         Returns an Umbral public key as bytes.
         Optionally, if an encoder function is provided it will be used to encode
@@ -266,7 +260,7 @@ class UmbralKeyingMaterial(object):
     
     """
 
-    def __init__(self, keying_material: bytes=None):
+    def __init__(self, keying_material: bytes = None):
         """
         Initializes an UmbralKeyingMaterial.
         """
@@ -277,8 +271,8 @@ def __init__(self, keying_material: bytes=None):
         else:
             self.keying_material = os.urandom(64)
 
-    def derive_privkey_by_label(self, label: bytes, salt: bytes=None, 
-                                params: UmbralParameters=None):
+    def derive_privkey_by_label(self, label: bytes, salt: bytes = None, 
+                                params: UmbralParameters = None):
         """
         Derives an UmbralPrivateKey using a KDF from this instance of 
         UmbralKeyingMaterial, a label, and an optional salt.
@@ -326,7 +320,7 @@ def from_bytes(cls, key_bytes: bytes, password: bytes=None, _scrypt_cost: int=20
 
         return cls(key_bytes)
 
-    def to_bytes(self, password: bytes=None, _scrypt_cost: int=20):
+    def to_bytes(self, password: bytes = None, _scrypt_cost: int = 20):
         """
         Returns an UmbralKeyingMaterial as a urlsafe base64 encoded string with
         optional symmetric encryption via nacl's Salsa20-Poly1305 and Scrypt
diff --git a/umbral/point.py b/umbral/point.py
index 1b0ce9a4..ba72648d 100644
--- a/umbral/point.py
+++ b/umbral/point.py
@@ -7,6 +7,7 @@
 from umbral import openssl
 from umbral.config import default_curve
 from umbral.utils import get_field_order_size_in_bytes
+from umbral.params import UmbralParameters
 
 
 class Point(object):
@@ -231,7 +232,7 @@ def __bytes__(self):
         return self.to_bytes()
 
 
-def unsafe_hash_to_point(data, params, label=None):
+def unsafe_hash_to_point(data, params : UmbralParameters, label=None):
     """
     Hashes arbitrary data into a valid EC point of the specified curve,
     using the try-and-increment method.
diff --git a/umbral/pre.py b/umbral/pre.py
index 5482849e..82ef4aeb 100644
--- a/umbral/pre.py
+++ b/umbral/pre.py
@@ -8,7 +8,7 @@
 from umbral._pre import prove_cfrag_correctness, assess_cfrag_correctness
 from umbral.curvebn import CurveBN
 from umbral.config import default_params, default_curve
-from umbral.dem import UmbralDEM
+from umbral.dem import UmbralDEM, DEM_KEYSIZE
 from umbral.fragments import KFrag, CapsuleFrag
 from umbral.keys import UmbralPrivateKey, UmbralPublicKey
 from umbral.params import UmbralParameters
@@ -17,8 +17,6 @@
 from umbral.utils import poly_eval, lambda_coeff, kdf
 import os
 
-CHACHA20_KEY_SIZE = 32
-
 
 class GenericUmbralError(Exception):
     pass
@@ -32,6 +30,7 @@ def __init__(self, message, offending_cfrags):
 
 class Capsule(object):
     def __init__(self,
+                 params: UmbralParameters,
                  point_e=None,
                  point_v=None,
                  bn_sig=None,
@@ -40,10 +39,9 @@ def __init__(self,
                  point_noninteractive=None,
                  delegating_pubkey: UmbralPublicKey = None,
                  receiving_pubkey: UmbralPublicKey = None,
-                 verifying_pubkey: UmbralPublicKey = None,
-                 params: UmbralParameters = None):
+                 verifying_pubkey = None):
 
-        self._umbral_params = params if params is not None else default_params()
+        self._umbral_params = params
 
         if isinstance(point_e, Point):
             if not isinstance(point_v, Point) or not isinstance(bn_sig, CurveBN):
@@ -91,11 +89,10 @@ class NotValid(ValueError):
         """
 
     @classmethod
-    def from_bytes(cls, capsule_bytes: bytes, params: UmbralParameters = None):
+    def from_bytes(cls, capsule_bytes: bytes, params: UmbralParameters):
         """
         Instantiates a Capsule object from the serialized data.
         """
-        params = params if params is not None else default_params()
         curve = params.curve
 
         bn_size = CurveBN.expected_bytes_length(curve)
@@ -121,7 +118,7 @@ def from_bytes(cls, capsule_bytes: bytes, params: UmbralParameters = None):
             raise ValueError("Byte string does not have a valid length for a Capsule")
 
         components = splitter(capsule_bytes)
-        return cls(*components, params=params)
+        return cls(params, *components)
 
     def _set_cfrag_correctness_key(self, key_type, key: UmbralPublicKey):
         if key_type not in ("delegating", "receiving", "verifying"): 
@@ -184,7 +181,6 @@ def verify_cfrag(self, cfrag):
                                         self._cfrag_correctness_keys["delegating"],
                                         self._cfrag_correctness_keys["receiving"],
                                         self._cfrag_correctness_keys["verifying"],
-                                        self._umbral_params
                                         )
 
     def original_components(self) -> Tuple[Point, Point, CurveBN]:
@@ -196,7 +192,8 @@ def activated_components(self) -> Union[Tuple[None, None, None], Tuple[Point, Po
     def _reconstruct_shamirs_secret(self,
                                     priv_b: Union[UmbralPrivateKey, CurveBN]
                                     ) -> None:
-        g = self._umbral_params.g
+        params = self._umbral_params
+        g = params.g
 
         if isinstance(priv_b, UmbralPrivateKey):
             pub_b = priv_b.get_pubkey()
@@ -218,9 +215,9 @@ def _reconstruct_shamirs_secret(self,
         hashed_dh_tuple = blake2b.finalize()
 
         if len(self._attached_cfrags) > 1:
-            xs = [CurveBN.hash(cfrag._kfrag_id, hashed_dh_tuple, params=self._umbral_params)
+            xs = [CurveBN.hash(cfrag._kfrag_id, hashed_dh_tuple, params=params)
                   for cfrag in self._attached_cfrags]
-            x_0 = CurveBN.hash(id_0, hashed_dh_tuple, params=self._umbral_params)
+            x_0 = CurveBN.hash(id_0, hashed_dh_tuple, params=params)
             lambda_0 = lambda_coeff(x_0, xs)
             e = lambda_0 * cfrag_0._point_e1
             v = lambda_0 * cfrag_0._point_v1
@@ -229,7 +226,7 @@ def _reconstruct_shamirs_secret(self,
                 if (ni, xcoord) != (cfrag._point_noninteractive, cfrag._point_xcoord):
                     raise ValueError("Attached CFrags are not pairwise consistent")
 
-                x_i = CurveBN.hash(cfrag._kfrag_id, hashed_dh_tuple, params=self._umbral_params)
+                x_i = CurveBN.hash(cfrag._kfrag_id, hashed_dh_tuple, params=params)
                 lambda_i = lambda_coeff(x_i, xs)
                 e = e + (lambda_i * cfrag._point_e1)
                 v = v + (lambda_i * cfrag._point_v1)
@@ -277,11 +274,9 @@ def __len__(self):
         return len(self._attached_cfrags)
 
 
-def split_rekey(privkey_a_bn: Union[UmbralPrivateKey, CurveBN],
-                signer_a: Signer,
-                pubkey_b_point: Union[UmbralPublicKey, Point],
-                threshold: int, N: int,
-                params: UmbralParameters = None) -> List[KFrag]:
+def split_rekey(delegating_privkey: UmbralPrivateKey, signer: Signer,
+                receiving_pubkey: UmbralPublicKey,
+                threshold: int, N: int) -> List[KFrag]:
     """
     Creates a re-encryption key from Alice to Bob and splits it in KFrags,
     using Shamir's Secret Sharing. Requires a threshold number of KFrags
@@ -289,18 +284,18 @@ def split_rekey(privkey_a_bn: Union[UmbralPrivateKey, CurveBN],
 
     Returns a list of KFrags.
     """
-    params = params if params is not None else default_params()
+
+    if delegating_privkey.params != receiving_pubkey.params:
+        raise ValueError("Keys must have the same parameter set.")
+
+    params = delegating_privkey.params
 
     g = params.g
 
-    if isinstance(privkey_a_bn, UmbralPrivateKey):
-        pubkey_a_point = privkey_a_bn.get_pubkey().point_key
-        privkey_a_bn = privkey_a_bn.bn_key
-    else:
-        pubkey_a_point = privkey_a_bn * g
+    pubkey_a_point = delegating_privkey.get_pubkey().point_key
+    privkey_a_bn = delegating_privkey.bn_key
 
-    if isinstance(pubkey_b_point, UmbralPublicKey):
-        pubkey_b_point = pubkey_b_point.point_key
+    pubkey_b_point = receiving_pubkey.point_key
 
     # 'ni' stands for 'Non Interactive'.
     # This point is used as an ephemeral public key in a DH key exchange,
@@ -343,7 +338,7 @@ def split_rekey(privkey_a_bn: Union[UmbralPrivateKey, CurveBN],
 
         kfrag_validity_message = bytes().join(
             bytes(material) for material in (id, pubkey_a_point, pubkey_b_point, u1, ni, xcoord))
-        signature = signer_a(kfrag_validity_message)
+        signature = signer(kfrag_validity_message)
 
         kfrag = KFrag(id=id, bn_key=rk,
                       point_noninteractive=ni, point_commitment=u1,
@@ -353,14 +348,13 @@ def split_rekey(privkey_a_bn: Union[UmbralPrivateKey, CurveBN],
     return kfrags
 
 
-def reencrypt(kfrag: KFrag, capsule: Capsule, params: UmbralParameters = None,
-              provide_proof=True, metadata: bytes = None) -> CapsuleFrag:
-    if params is None:
-        params = default_params()
+def reencrypt(kfrag: KFrag, capsule: Capsule, provide_proof = True, 
+              metadata: bytes = None) -> CapsuleFrag:
 
     if not capsule.verify():
         raise capsule.NotValid
 
+
     rk = kfrag._bn_key
     e1 = rk * capsule._point_e
     v1 = rk * capsule._point_v
@@ -370,15 +364,14 @@ def reencrypt(kfrag: KFrag, capsule: Capsule, params: UmbralParameters = None,
                         point_xcoord=kfrag._point_xcoord)
 
     if provide_proof:
-        prove_cfrag_correctness(cfrag, kfrag, capsule, metadata, params)
+        prove_cfrag_correctness(cfrag, kfrag, capsule, metadata)
 
     return cfrag
 
 
-def _encapsulate(alice_pub_key: Point, key_length=32,
-                 params: UmbralParameters = None) -> Tuple[bytes, Capsule]:
+def _encapsulate(alice_pubkey: Point, params: UmbralParameters,
+                 key_length = DEM_KEYSIZE) -> Tuple[bytes, Capsule]:
     """Generates a symmetric key and its associated KEM ciphertext"""
-    params = params if params is not None else default_params()
 
     g = params.g
 
@@ -391,7 +384,7 @@ def _encapsulate(alice_pub_key: Point, key_length=32,
     h = CurveBN.hash(pub_r, pub_u, params=params)
     s = priv_u + (priv_r * h)
 
-    shared_key = (priv_r + priv_u) * alice_pub_key
+    shared_key = (priv_r + priv_u) * alice_pubkey
 
     # Key to be used for symmetric encryption
     key = kdf(shared_key, key_length)
@@ -399,10 +392,9 @@ def _encapsulate(alice_pub_key: Point, key_length=32,
     return key, Capsule(point_e=pub_r, point_v=pub_u, bn_sig=s, params=params)
 
 
-def _decapsulate_original(priv_key: CurveBN, capsule: Capsule, key_length=32,
-                          params: UmbralParameters = None) -> bytes:
+def _decapsulate_original(priv_key: CurveBN, capsule: Capsule, 
+                          key_length = DEM_KEYSIZE) -> bytes:
     """Derive the same symmetric key"""
-    params = params if params is not None else default_params()
 
     shared_key = priv_key * (capsule._point_e + capsule._point_v)
     key = kdf(shared_key, key_length)
@@ -417,9 +409,9 @@ def _decapsulate_original(priv_key: CurveBN, capsule: Capsule, key_length=32,
 
 def _decapsulate_reencrypted(pub_key: Point, priv_key: CurveBN,
                              orig_pub_key: Point, capsule: Capsule,
-                             key_length=32, params: UmbralParameters = None) -> bytes:
+                             key_length = DEM_KEYSIZE) -> bytes:
     """Derive the same symmetric key"""
-    params = params if params is not None else default_params()
+    params = capsule._umbral_params
 
     ni = capsule._point_noninteractive
     d = CurveBN.hash(ni, pub_key, priv_key * ni, params=params)
@@ -442,17 +434,16 @@ def _decapsulate_reencrypted(pub_key: Point, priv_key: CurveBN,
     return key
 
 
-def encrypt(alice_pubkey: UmbralPublicKey, plaintext: bytes,
-            params: UmbralParameters = None) -> Tuple[bytes, Capsule]:
+def encrypt(alice_pubkey: UmbralPublicKey, plaintext: bytes) -> Tuple[bytes, Capsule]:
     """
     Performs an encryption using the UmbralDEM object and encapsulates a key
     for the sender using the public key provided.
 
     Returns the ciphertext and the KEM Capsule.
     """
-    params = params if params is not None else default_params()
+    params = alice_pubkey.params
 
-    key, capsule = _encapsulate(alice_pubkey.point_key, CHACHA20_KEY_SIZE, params=params)
+    key, capsule = _encapsulate(alice_pubkey.point_key, params, DEM_KEYSIZE)
 
     capsule_bytes = bytes(capsule)
 
@@ -463,10 +454,9 @@ def encrypt(alice_pubkey: UmbralPublicKey, plaintext: bytes,
 
 
 def _open_capsule(capsule: Capsule,
-                  bob_privkey: UmbralPrivateKey,
+                  receiving_privkey: UmbralPrivateKey,
                   delegating_pubkey: UmbralPublicKey,
-                  alice_pubkey: UmbralPublicKey,
-                  params: UmbralParameters = None,
+                  signing_pubkey: UmbralPublicKey,
                   check_proof=True) -> bytes:
     """
     Activates the Capsule from the attached CFrags,
@@ -474,19 +464,17 @@ def _open_capsule(capsule: Capsule,
 
     This will often be a symmetric key.
     """
-    params = params if params is not None else default_params()
 
-    priv_b = bob_privkey.bn_key
-    bob_pubkey = bob_privkey.get_pubkey()
+    receiving_pubkey = receiving_privkey.get_pubkey()
+    priv_b = receiving_privkey.bn_key
 
     if check_proof:
         offending_cfrags = []
         for cfrag in capsule._attached_cfrags:
             if not cfrag.verify_correctness(capsule=capsule,
                                             delegating_pubkey=delegating_pubkey,
-                                            signing_pubkey=alice_pubkey,
-                                            receiving_pubkey=bob_pubkey,
-                                            params=params):
+                                            signing_pubkey=signing_pubkey,
+                                            receiving_pubkey=receiving_pubkey):
                 offending_cfrags.append(cfrag)
 
         if offending_cfrags:
@@ -495,7 +483,8 @@ def _open_capsule(capsule: Capsule,
 
     capsule._reconstruct_shamirs_secret(priv_b)
 
-    key = _decapsulate_reencrypted(bob_pubkey.point_key, priv_b, delegating_pubkey.point_key, capsule, params=params)
+    key = _decapsulate_reencrypted(receiving_pubkey.point_key, priv_b, 
+                                   delegating_pubkey.point_key, capsule)
     return key
 
 
@@ -503,22 +492,22 @@ def decrypt(ciphertext: bytes,
             capsule: Capsule,
             decrypting_key: UmbralPrivateKey,
             delegating_pubkey: UmbralPublicKey = None,
-            verifying_key: UmbralPublicKey = None,
-            params: UmbralParameters = None, check_proof=True) -> bytes:
+            verifying_key = None,
+            check_proof=True) -> bytes:
     """
     Opens the capsule and gets what's inside.
 
     We hope that's a symmetric key, which we use to decrypt the ciphertext
     and return the resulting cleartext.
     """
-    params = params if params is not None else default_params()
 
     if capsule._attached_cfrags:
         # Since there are cfrags attached, we assume this is Bob opening the Capsule.
         # (i.e., this is a re-encrypted capsule)
 
-        encapsulated_key = _open_capsule(capsule, decrypting_key, delegating_pubkey, verifying_key,
-                                         params=params, check_proof=check_proof)
+        encapsulated_key = _open_capsule(capsule, decrypting_key, 
+                                         delegating_pubkey, verifying_key,
+                                         check_proof=check_proof)
         dem = UmbralDEM(encapsulated_key)
 
         original_capsule_bytes = capsule._original_to_bytes()
@@ -526,7 +515,7 @@ def decrypt(ciphertext: bytes,
     else:
         # Since there aren't cfrags attached, we assume this is Alice opening the Capsule.
         # (i.e., this is an original capsule)
-        decapsulated_key = _decapsulate_original(decrypting_key.bn_key, capsule, params=params)
+        decapsulated_key = _decapsulate_original(decrypting_key.bn_key, capsule)
         dem = UmbralDEM(decapsulated_key)
 
         capsule_bytes = bytes(capsule)

From d5b434c0130f83dc722f8277e5765db5292803ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Fri, 1 Jun 2018 12:30:19 +0200
Subject: [PATCH 06/11] Adapts tests to new way to use params

---
 .../test_capsule_correctness_checks.py        | 17 +++++--
 tests/test_capsule/test_capsule_operations.py | 30 ++++++++---
 .../test_capsule/test_capsule_serializers.py  | 27 ++++++----
 tests/test_correctness.py                     | 51 ++++++++++++-------
 tests/test_keys/test_key_fragments.py         |  6 +--
 tests/test_keys/test_umbral_keys.py           | 10 ++--
 tests/test_simple_api.py                      | 11 ++--
 7 files changed, 100 insertions(+), 52 deletions(-)

diff --git a/tests/test_capsule/test_capsule_correctness_checks.py b/tests/test_capsule/test_capsule_correctness_checks.py
index bac12cf3..c8d5f9c7 100644
--- a/tests/test_capsule/test_capsule_correctness_checks.py
+++ b/tests/test_capsule/test_capsule_correctness_checks.py
@@ -9,6 +9,7 @@
 from umbral.point import Point
 from umbral.pre import Capsule
 from umbral.signing import Signer
+from umbral.config import default_params
 
 
 def test_cannot_attach_cfrag_without_keys():
@@ -16,7 +17,11 @@ def test_cannot_attach_cfrag_without_keys():
     We need the proper keys to verify the correctness of CFrags
     in order to attach them to a Capsule.
     """
-    capsule = Capsule(point_e=Point.gen_rand(),
+
+    params = default_params()
+
+    capsule = Capsule(params,
+                      point_e=Point.gen_rand(),
                       point_v=Point.gen_rand(),
                       bn_sig=CurveBN.gen_rand())
 
@@ -59,7 +64,10 @@ def test_cannot_attach_cfrag_without_proof():
     However, even when properly attaching keys, we can't attach the CFrag
     if it is unproven.
     """
-    capsule = Capsule(point_e=Point.gen_rand(),
+    params = default_params()
+
+    capsule = Capsule(params,
+                      point_e=Point.gen_rand(),
                       point_v=Point.gen_rand(),
                       bn_sig=CurveBN.gen_rand())
 
@@ -86,7 +94,10 @@ def test_cannot_set_different_keys():
     """
     Once a key is set on a Capsule, it can't be changed to a different key.
     """
-    capsule = Capsule(point_e=Point.gen_rand(),
+    params = default_params()
+
+    capsule = Capsule(params,
+                      point_e=Point.gen_rand(),
                       point_v=Point.gen_rand(),
                       bn_sig=CurveBN.gen_rand())
 
diff --git a/tests/test_capsule/test_capsule_operations.py b/tests/test_capsule/test_capsule_operations.py
index 764719cc..a52c9cf8 100644
--- a/tests/test_capsule/test_capsule_operations.py
+++ b/tests/test_capsule/test_capsule_operations.py
@@ -6,14 +6,21 @@
 from umbral.pre import Capsule
 from umbral.signing import Signer
 from umbral.keys import UmbralPrivateKey
+from umbral.config import default_params
 
 
 def test_capsule_creation(alices_keys):
+
+    params = default_params()
+
     with pytest.raises(TypeError):
-        rare_capsule = Capsule()  # Alice cannot make a capsule this way.
+        rare_capsule = Capsule(params)  # Alice cannot make a capsule this way.
+
+
 
     # Some users may create capsules their own way.
-    custom_capsule = Capsule(point_e=Point.gen_rand(),
+    custom_capsule = Capsule(params,
+                             point_e=Point.gen_rand(),
                              point_v=Point.gen_rand(),
                              bn_sig=CurveBN.gen_rand())
 
@@ -28,17 +35,22 @@ def test_capsule_creation(alices_keys):
 
 
 def test_capsule_equality():
-    one_capsule = Capsule(point_e=Point.gen_rand(),
+    params = default_params()
+
+    one_capsule = Capsule(params,
+                          point_e=Point.gen_rand(),
                           point_v=Point.gen_rand(),
                           bn_sig=CurveBN.gen_rand())
 
-    another_capsule = Capsule(point_e=Point.gen_rand(),
+    another_capsule = Capsule(params,
+                              point_e=Point.gen_rand(),
                               point_v=Point.gen_rand(),
                               bn_sig=CurveBN.gen_rand())
 
     assert one_capsule != another_capsule
 
-    activated_capsule = Capsule(point_e_prime=Point.gen_rand(),
+    activated_capsule = Capsule(params,
+                                point_e_prime=Point.gen_rand(),
                                 point_v_prime=Point.gen_rand(),
                                 point_noninteractive=Point.gen_rand())
 
@@ -46,9 +58,11 @@ def test_capsule_equality():
 
 
 def test_decapsulation_by_alice(alices_keys):
+    params = default_params()
+
     delegating_privkey, _signing_privkey = alices_keys
 
-    sym_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey().point_key)
+    sym_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey().point_key, params)
     assert len(sym_key) == 32
 
     # The symmetric key sym_key is perhaps used for block cipher here in a real-world scenario.
@@ -57,6 +71,7 @@ def test_decapsulation_by_alice(alices_keys):
 
 
 def test_bad_capsule_fails_reencryption(alices_keys, bobs_keys):
+    params = default_params()
     delegating_privkey, _signing_privkey = alices_keys
     signer_alice = Signer(_signing_privkey)
 
@@ -64,7 +79,8 @@ def test_bad_capsule_fails_reencryption(alices_keys, bobs_keys):
 
     kfrags = pre.split_rekey(delegating_privkey, signer_alice, receiving_pubkey, 1, 2)
 
-    bollocks_capsule = Capsule(point_e=Point.gen_rand(),
+    bollocks_capsule = Capsule(params,
+                               point_e=Point.gen_rand(),
                                point_v=Point.gen_rand(),
                                bn_sig=CurveBN.gen_rand())
 
diff --git a/tests/test_capsule/test_capsule_serializers.py b/tests/test_capsule/test_capsule_serializers.py
index bb93ec7d..008efe22 100644
--- a/tests/test_capsule/test_capsule_serializers.py
+++ b/tests/test_capsule/test_capsule_serializers.py
@@ -3,14 +3,15 @@
 from umbral import pre, keys
 from umbral.curvebn import CurveBN
 from umbral.point import Point
-from umbral.config import default_curve
+from umbral.config import default_curve, default_params
 from umbral.signing import Signer
 
 
 def test_capsule_serialization(alices_keys):
     delegating_privkey, _signing_privkey = alices_keys
+    params = delegating_privkey.params
 
-    _symmetric_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey().point_key)
+    _symmetric_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey().point_key, params)
     capsule_bytes = capsule.to_bytes()
     capsule_bytes_casted = bytes(capsule)
     assert capsule_bytes == capsule_bytes_casted
@@ -18,7 +19,7 @@ def test_capsule_serialization(alices_keys):
     # A Capsule can be represented as the 98 total bytes of two Points (33 each) and a CurveBN (32).
     assert len(capsule_bytes) == pre.Capsule.expected_bytes_length()
 
-    new_capsule = pre.Capsule.from_bytes(capsule_bytes)
+    new_capsule = pre.Capsule.from_bytes(capsule_bytes, params)
 
     # Three ways to think about equality.
     # First, the public approach for the Capsule.  Simply:
@@ -39,9 +40,11 @@ def test_activated_capsule_serialization(alices_keys, bobs_keys):
     delegating_pubkey = delegating_privkey.get_pubkey()
     signer_alice = Signer(signing_privkey)
 
+    params = delegating_privkey.params
+
     receiving_privkey, receiving_pubkey = bobs_keys
 
-    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key)
+    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key, params)
 
     capsule.set_correctness_keys(delegating=delegating_pubkey,
                                  receiving=receiving_pubkey,
@@ -58,7 +61,7 @@ def test_activated_capsule_serialization(alices_keys, bobs_keys):
 
     assert len(rec_capsule_bytes) == pre.Capsule.expected_bytes_length(activated=True)
 
-    new_rec_capsule = pre.Capsule.from_bytes(rec_capsule_bytes)
+    new_rec_capsule = pre.Capsule.from_bytes(rec_capsule_bytes, params)
 
     # Again, the same three perspectives on equality.
     assert new_rec_capsule == capsule
@@ -71,22 +74,28 @@ def test_activated_capsule_serialization(alices_keys, bobs_keys):
 
 
 def test_cannot_create_capsule_from_bogus_material(alices_keys):
+    params = alices_keys[0].params
+    
     with pytest.raises(TypeError):
-        capsule_of_questionable_parentage = pre.Capsule(point_e=Point.gen_rand(),
+        capsule_of_questionable_parentage = pre.Capsule(params,
+                                                        point_e=Point.gen_rand(),
                                                         point_v=42,
                                                         bn_sig=CurveBN.gen_rand())
 
     with pytest.raises(TypeError):
-        capsule_of_questionable_parentage = pre.Capsule(point_e=Point.gen_rand(),
+        capsule_of_questionable_parentage = pre.Capsule(params,
+                                                        point_e=Point.gen_rand(),
                                                         point_v=Point.gen_rand(),
                                                         bn_sig=42)
 
     with pytest.raises(TypeError):
-        capsule_of_questionable_parentage = pre.Capsule(point_e_prime=Point.gen_rand(),
+        capsule_of_questionable_parentage = pre.Capsule(params,
+                                                        point_e_prime=Point.gen_rand(),
                                                         point_v_prime=42,
                                                         point_noninteractive=Point.gen_rand())
 
     with pytest.raises(TypeError):
-        capsule_of_questionable_parentage = pre.Capsule(point_e_prime=Point.gen_rand(),
+        capsule_of_questionable_parentage = pre.Capsule(params,
+                                                        point_e_prime=Point.gen_rand(),
                                                         point_v_prime=Point.gen_rand(),
                                                         point_noninteractive=42)
diff --git a/tests/test_correctness.py b/tests/test_correctness.py
index 98ae24d9..7e985dbc 100644
--- a/tests/test_correctness.py
+++ b/tests/test_correctness.py
@@ -9,12 +9,15 @@
 
 def test_correctness_proof_serialization(alices_keys):
     delegating_privkey, signing_privkey = alices_keys
+    delegating_pubkey = delegating_privkey.get_pubkey()
     signer = Signer(signing_privkey)
 
     priv_key_bob = keys.UmbralPrivateKey.gen_key()
     pub_key_bob = priv_key_bob.get_pubkey()
 
-    _unused_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey().point_key)
+    params = delegating_privkey.params
+
+    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key, params)
     kfrags = pre.split_rekey(delegating_privkey, signer, pub_key_bob, 1, 2)
 
     # Example of potential metadata to describe the re-encryption request
@@ -42,15 +45,18 @@ def test_correctness_proof_serialization(alices_keys):
 @pytest.mark.parametrize("N, M", parameters)
 def test_cheating_ursula_replays_old_reencryption(N, M, alices_keys):
     delegating_privkey, signing_privkey = alices_keys
+    delegating_pubkey = delegating_privkey.get_pubkey()
     signer = Signer(signing_privkey)
 
     priv_key_bob = keys.UmbralPrivateKey.gen_key()
     pub_key_bob = priv_key_bob.get_pubkey()
 
-    sym_key_alice1, capsule_alice1 = pre._encapsulate(delegating_privkey.get_pubkey().point_key)
-    sym_key_alice2, capsule_alice2 = pre._encapsulate(delegating_privkey.get_pubkey().point_key)
+    params = delegating_privkey.params
+
+    _unused_key1, capsule_alice1 = pre._encapsulate(delegating_pubkey.point_key, params)
+    _unused_key2, capsule_alice2 = pre._encapsulate(delegating_pubkey.point_key, params)
 
-    capsule_alice1.set_correctness_keys(delegating=delegating_privkey.get_pubkey(),
+    capsule_alice1.set_correctness_keys(delegating=delegating_pubkey,
                                         receiving=pub_key_bob,
                                         verifying=signing_privkey.get_pubkey())
 
@@ -79,12 +85,12 @@ def test_cheating_ursula_replays_old_reencryption(N, M, alices_keys):
     with pytest.raises(pre.GenericUmbralError):
         sym_key = pre._decapsulate_reencrypted(pub_key_bob.point_key,
                                                priv_key_bob.bn_key,
-                                               delegating_privkey.get_pubkey().point_key,
+                                               delegating_pubkey.point_key,
                                                capsule_alice1
                                                )
 
     assert not cfrags[0].verify_correctness(capsule_alice1,
-                                            delegating_privkey.get_pubkey(),
+                                            delegating_pubkey,
                                             signing_privkey.get_pubkey(),
                                             pub_key_bob,
                                             )
@@ -94,7 +100,7 @@ def test_cheating_ursula_replays_old_reencryption(N, M, alices_keys):
     correct_cases = 0
     for cfrag_i, metadata_i in zip(cfrags[1:], metadata[1:]):
         if cfrag_i.verify_correctness(capsule_alice1,
-                                      delegating_privkey.get_pubkey(),
+                                      delegating_pubkey,
                                       signing_privkey.get_pubkey(),
                                       pub_key_bob,
                                       ):
@@ -108,7 +114,7 @@ def test_cheating_ursula_replays_old_reencryption(N, M, alices_keys):
     with pytest.raises(pre.UmbralCorrectnessError) as exception_info:
         _ = pre._open_capsule(capsule_alice1,
                               priv_key_bob,
-                              delegating_privkey.get_pubkey(),
+                              delegating_pubkey,
                               signing_privkey.get_pubkey(),
                               )
     correctness_error = exception_info.value
@@ -119,15 +125,17 @@ def test_cheating_ursula_replays_old_reencryption(N, M, alices_keys):
 @pytest.mark.parametrize("N, M", parameters)
 def test_cheating_ursula_sends_garbage(N, M, alices_keys):
     delegating_privkey, signing_privkey = alices_keys
+    delegating_pubkey = delegating_privkey.get_pubkey()
     signer = Signer(signing_privkey)
 
-    # Bob
     priv_key_bob = keys.UmbralPrivateKey.gen_key()
     pub_key_bob = priv_key_bob.get_pubkey()
 
-    sym_key, capsule_alice = pre._encapsulate(delegating_privkey.get_pubkey().point_key)
+    params = delegating_privkey.params
 
-    capsule_alice.set_correctness_keys(delegating=delegating_privkey.get_pubkey(),
+    sym_key, capsule_alice = pre._encapsulate(delegating_pubkey.point_key, params)
+
+    capsule_alice.set_correctness_keys(delegating=delegating_pubkey,
                                        receiving=pub_key_bob,
                                        verifying=signing_privkey.get_pubkey())
 
@@ -153,11 +161,11 @@ def test_cheating_ursula_sends_garbage(N, M, alices_keys):
     with pytest.raises(pre.GenericUmbralError):
         _unused_key = pre._decapsulate_reencrypted(pub_key_bob.point_key,
                                                    priv_key_bob.bn_key,
-                                                   delegating_privkey.get_pubkey().point_key,
+                                                   delegating_pubkey.point_key,
                                                    capsule_alice)
 
     assert not cfrags[0].verify_correctness(capsule_alice,
-                                            delegating_privkey.get_pubkey(),
+                                            delegating_pubkey,
                                             signing_privkey.get_pubkey(),
                                             pub_key_bob,
                                             )
@@ -166,7 +174,7 @@ def test_cheating_ursula_sends_garbage(N, M, alices_keys):
     # so the rest of CFrags chould be correct:
     for cfrag_i, metadata_i in zip(cfrags[1:], metadata[1:]):
         assert cfrag_i.verify_correctness(capsule_alice,
-                                          delegating_privkey.get_pubkey(),
+                                          delegating_pubkey,
                                           signing_privkey.get_pubkey(),
                                           pub_key_bob,
                                           )
@@ -174,7 +182,8 @@ def test_cheating_ursula_sends_garbage(N, M, alices_keys):
     # Alternatively, we can try to open the capsule directly.
     # We should get an exception with an attached list of incorrect cfrags
     with pytest.raises(pre.UmbralCorrectnessError) as exception_info:
-        _decapsulated_key = pre._open_capsule(capsule_alice, priv_key_bob, delegating_privkey.get_pubkey(),
+        _decapsulated_key = pre._open_capsule(capsule_alice, priv_key_bob, 
+                                              delegating_pubkey,
                                               signing_privkey.get_pubkey())
     correctness_error = exception_info.value
     assert cfrags[0] in correctness_error.offending_cfrags
@@ -186,6 +195,7 @@ def test_decryption_fails_when_it_expects_a_proof_and_there_isnt(N, M, alices_ke
     """Manually injects umbralparameters for multi-curve testing."""
 
     delegating_privkey, signing_privkey = alices_keys
+    delegating_pubkey = delegating_privkey.get_pubkey()
     signer = Signer(signing_privkey)
     priv_key_bob, pub_key_bob = bobs_keys
 
@@ -215,11 +225,14 @@ def test_decryption_fails_when_it_expects_a_proof_and_there_isnt(N, M, alices_ke
 @pytest.mark.parametrize("N, M", parameters)
 def test_m_of_n(N, M, alices_keys, bobs_keys):
     delegating_privkey, signing_privkey = alices_keys
+    delegating_pubkey = delegating_privkey.get_pubkey()
     signer = Signer(signing_privkey)
 
     priv_key_bob, pub_key_bob = bobs_keys
 
-    sym_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey().point_key)
+    params = delegating_privkey.params
+
+    sym_key, capsule = pre._encapsulate(delegating_pubkey.point_key, params)
 
     capsule.set_correctness_keys(delegating=delegating_privkey.get_pubkey(),
                                  receiving=pub_key_bob,
@@ -238,10 +251,10 @@ def test_m_of_n(N, M, alices_keys, bobs_keys):
         cfrag = pre.reencrypt(kfrag, capsule, metadata=metadata)
         capsule.attach_cfrag(cfrag)
 
-        assert cfrag.verify_correctness(capsule,
-                                        delegating_privkey.get_pubkey(), signing_privkey.get_pubkey(), pub_key_bob,
+        assert cfrag.verify_correctness(capsule, delegating_pubkey, 
+                                        signing_privkey.get_pubkey(), pub_key_bob,
                                         )
 
-    sym_key_from_capsule = pre._open_capsule(capsule, priv_key_bob, delegating_privkey.get_pubkey(),
+    sym_key_from_capsule = pre._open_capsule(capsule, priv_key_bob, delegating_pubkey,
                                              signing_privkey.get_pubkey())
     assert sym_key == sym_key_from_capsule
diff --git a/tests/test_keys/test_key_fragments.py b/tests/test_keys/test_key_fragments.py
index 8b145b0a..f25a6eab 100644
--- a/tests/test_keys/test_key_fragments.py
+++ b/tests/test_keys/test_key_fragments.py
@@ -32,7 +32,7 @@ def test_cfrag_serialization_with_proof_and_metadata(alices_keys, bobs_keys):
     _receiving_privkey, receiving_pubkey = bobs_keys
 
 
-    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key)
+    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key, delegating_pubkey.params)
     kfrags = pre.split_rekey(delegating_privkey, signer_alice,
                              receiving_pubkey, 1, 2)
 
@@ -70,7 +70,7 @@ def test_cfrag_serialization_with_proof_but_no_metadata(alices_keys, bobs_keys):
     _receiving_privkey, receiving_pubkey = bobs_keys
     signer_alice = Signer(signing_privkey)
 
-    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key)
+    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key, delegating_pubkey.params)
     kfrags = pre.split_rekey(delegating_privkey, signer_alice,
                              receiving_pubkey, 1, 2)
 
@@ -108,7 +108,7 @@ def test_cfrag_serialization_no_proof_no_metadata(alices_keys, bobs_keys):
     _receiving_privkey, receiving_pubkey = bobs_keys
     signer_alice = Signer(signing_privkey)
 
-    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key)
+    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key, delegating_pubkey.params)
     kfrags = pre.split_rekey(delegating_privkey, signer_alice,
                              receiving_pubkey, 1, 2)
 
diff --git a/tests/test_keys/test_umbral_keys.py b/tests/test_keys/test_umbral_keys.py
index df0783f3..be983e6b 100644
--- a/tests/test_keys/test_umbral_keys.py
+++ b/tests/test_keys/test_umbral_keys.py
@@ -46,7 +46,7 @@ def test_derive_key_from_label():
 
 def test_private_key_serialization(random_ec_curvebn1):
     priv_key = random_ec_curvebn1
-    umbral_key = keys.UmbralPrivateKey(priv_key)
+    umbral_key = keys.UmbralPrivateKey(priv_key, default_params())
 
     encoded_key = umbral_key.to_bytes()
 
@@ -56,7 +56,7 @@ def test_private_key_serialization(random_ec_curvebn1):
 
 def test_private_key_serialization_with_encryption(random_ec_curvebn1):
     priv_key = random_ec_curvebn1
-    umbral_key = keys.UmbralPrivateKey(priv_key)
+    umbral_key = keys.UmbralPrivateKey(priv_key, default_params())
 
     encoded_key = umbral_key.to_bytes(password=b'test')
 
@@ -70,7 +70,7 @@ def test_public_key_serialization(random_ec_curvebn1):
     params = default_params()
     pub_key = priv_key * params.g
 
-    umbral_key = keys.UmbralPublicKey(pub_key)
+    umbral_key = keys.UmbralPublicKey(pub_key, params)
 
     encoded_key = umbral_key.to_bytes()
 
@@ -84,7 +84,7 @@ def test_public_key_to_bytes(random_ec_curvebn1):
     params = default_params()
     pub_key = priv_key * params.g
 
-    umbral_key = keys.UmbralPublicKey(pub_key)
+    umbral_key = keys.UmbralPublicKey(pub_key, params)
     key_bytes = bytes(umbral_key)
 
     assert type(key_bytes) == bytes
@@ -92,7 +92,7 @@ def test_public_key_to_bytes(random_ec_curvebn1):
 
 def test_key_encoder_decoder(random_ec_curvebn1):
     priv_key = random_ec_curvebn1
-    umbral_key = keys.UmbralPrivateKey(priv_key)
+    umbral_key = keys.UmbralPrivateKey(priv_key, default_params())
 
     encoded_key = umbral_key.to_bytes(encoder=base64.urlsafe_b64encode)
 
diff --git a/tests/test_simple_api.py b/tests/test_simple_api.py
index 605a8369..0af43c12 100644
--- a/tests/test_simple_api.py
+++ b/tests/test_simple_api.py
@@ -31,24 +31,23 @@ def test_simple_api(N, M, curve=default_curve()):
     receiving_pubkey = receiving_privkey.get_pubkey()
 
     plain_data = b'peace at dawn'
-    ciphertext, capsule = pre.encrypt(delegating_pubkey, plain_data, params=params)
+    ciphertext, capsule = pre.encrypt(delegating_pubkey, plain_data)
 
-    cleartext = pre.decrypt(ciphertext, capsule, delegating_privkey, params=params)
+    cleartext = pre.decrypt(ciphertext, capsule, delegating_privkey)
     assert cleartext == plain_data
 
     capsule.set_correctness_keys(delegating=delegating_pubkey,
                                  receiving=receiving_pubkey,
                                  verifying=signing_pubkey)
 
-    kfrags = pre.split_rekey(delegating_privkey, signer, receiving_pubkey, M, N, params=params)
+    kfrags = pre.split_rekey(delegating_privkey, signer, receiving_pubkey, M, N)
 
     for kfrag in kfrags:
-        cfrag = pre.reencrypt(kfrag, capsule, params=params)
+        cfrag = pre.reencrypt(kfrag, capsule)
         capsule.attach_cfrag(cfrag)
 
     reenc_cleartext = pre.decrypt(ciphertext, capsule, receiving_privkey,
-                                  delegating_pubkey, signing_pubkey,
-                                  params=params)
+                                  delegating_pubkey, signing_pubkey)
     assert reenc_cleartext == plain_data
 
 

From 5618b2529289af227e1d6d68d873fbc84ec44b9d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Mon, 4 Jun 2018 14:37:32 +0200
Subject: [PATCH 07/11] Simplified public API for umbral.pre

* Functions in `pre` now only take Umbral keys as arguments, rather than primitive types (Point, CurveBN)
* Remove unnecessary arguments from public facing and internal methods when they can be extracted from a Capsule, UmbralPublicKey or UmbralPrivateKey
* Adds a getter in Capsule for correctness keys
* Adapts the test suite to new simplified API
---
 tests/test_capsule/test_capsule_operations.py | 13 ++-
 .../test_capsule/test_capsule_serializers.py  | 13 +--
 tests/test_correctness.py                     | 88 ++++++-------------
 tests/test_keys/test_key_fragments.py         |  6 +-
 tests/test_simple_api.py                      |  3 +-
 umbral/_pre.py                                | 18 +++-
 umbral/fragments.py                           | 15 +---
 umbral/pre.py                                 | 56 +++++-------
 8 files changed, 86 insertions(+), 126 deletions(-)

diff --git a/tests/test_capsule/test_capsule_operations.py b/tests/test_capsule/test_capsule_operations.py
index a52c9cf8..fb92df27 100644
--- a/tests/test_capsule/test_capsule_operations.py
+++ b/tests/test_capsule/test_capsule_operations.py
@@ -62,11 +62,11 @@ def test_decapsulation_by_alice(alices_keys):
 
     delegating_privkey, _signing_privkey = alices_keys
 
-    sym_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey().point_key, params)
+    sym_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey())
     assert len(sym_key) == 32
 
     # The symmetric key sym_key is perhaps used for block cipher here in a real-world scenario.
-    sym_key_2 = pre._decapsulate_original(delegating_privkey.bn_key, capsule)
+    sym_key_2 = pre._decapsulate_original(delegating_privkey, capsule)
     assert sym_key_2 == sym_key
 
 
@@ -115,8 +115,7 @@ def test_capsule_as_dict_key(alices_keys, bobs_keys):
     capsule.attach_cfrag(cfrag)
 
     # Even if we activate the capsule, it still serves as the same key.
-    cleartext = pre.decrypt(ciphertext, capsule, receiving_privkey,
-                            delegating_pubkey, signing_pubkey)
+    cleartext = pre.decrypt(ciphertext, capsule, receiving_privkey)
     assert some_dict[capsule] == "Thing that Bob wants to try per-Capsule"
     assert cleartext == plain_data
 
@@ -132,14 +131,14 @@ def test_capsule_length(alices_keys, bobs_keys):
 
     priv_key_bob, pub_key_bob = bobs_keys
 
-    sym_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey().point_key)
+    sym_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey())
+
+    kfrags = pre.split_rekey(delegating_privkey, signer, pub_key_bob, 10, 15)
 
     capsule.set_correctness_keys(delegating=delegating_privkey.get_pubkey(),
                                  receiving=pub_key_bob,
                                  verifying=signing_privkey.get_pubkey())
 
-    kfrags = pre.split_rekey(delegating_privkey, signer, pub_key_bob, 10, 15)
-
     for counter, kfrag in enumerate(kfrags):
         assert len(capsule) == counter
         cfrag = pre.reencrypt(kfrag, capsule)
diff --git a/tests/test_capsule/test_capsule_serializers.py b/tests/test_capsule/test_capsule_serializers.py
index 008efe22..8b104d9d 100644
--- a/tests/test_capsule/test_capsule_serializers.py
+++ b/tests/test_capsule/test_capsule_serializers.py
@@ -11,7 +11,7 @@ def test_capsule_serialization(alices_keys):
     delegating_privkey, _signing_privkey = alices_keys
     params = delegating_privkey.params
 
-    _symmetric_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey().point_key, params)
+    _symmetric_key, capsule = pre._encapsulate(delegating_privkey.get_pubkey())
     capsule_bytes = capsule.to_bytes()
     capsule_bytes_casted = bytes(capsule)
     assert capsule_bytes == capsule_bytes_casted
@@ -44,16 +44,17 @@ def test_activated_capsule_serialization(alices_keys, bobs_keys):
 
     receiving_privkey, receiving_pubkey = bobs_keys
 
-    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key, params)
-
-    capsule.set_correctness_keys(delegating=delegating_pubkey,
-                                 receiving=receiving_pubkey,
-                                 verifying=signing_privkey.get_pubkey())
+    _unused_key, capsule = pre._encapsulate(delegating_pubkey)
 
+    
     kfrags = pre.split_rekey(delegating_privkey, signer_alice, receiving_pubkey, 1, 2)
 
     cfrag = pre.reencrypt(kfrags[0], capsule)
 
+    capsule.set_correctness_keys(delegating=delegating_pubkey,
+                                 receiving=receiving_pubkey,
+                                 verifying=signing_privkey.get_pubkey())
+    
     capsule.attach_cfrag(cfrag)
 
     capsule._reconstruct_shamirs_secret(receiving_privkey)
diff --git a/tests/test_correctness.py b/tests/test_correctness.py
index 7e985dbc..1d25bd47 100644
--- a/tests/test_correctness.py
+++ b/tests/test_correctness.py
@@ -17,7 +17,7 @@ def test_correctness_proof_serialization(alices_keys):
 
     params = delegating_privkey.params
 
-    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key, params)
+    _unused_key, capsule = pre._encapsulate(delegating_pubkey)
     kfrags = pre.split_rekey(delegating_privkey, signer, pub_key_bob, 1, 2)
 
     # Example of potential metadata to describe the re-encryption request
@@ -53,15 +53,15 @@ def test_cheating_ursula_replays_old_reencryption(N, M, alices_keys):
 
     params = delegating_privkey.params
 
-    _unused_key1, capsule_alice1 = pre._encapsulate(delegating_pubkey.point_key, params)
-    _unused_key2, capsule_alice2 = pre._encapsulate(delegating_pubkey.point_key, params)
+    _unused_key1, capsule_alice1 = pre._encapsulate(delegating_pubkey)
+    _unused_key2, capsule_alice2 = pre._encapsulate(delegating_pubkey)
+
+    kfrags = pre.split_rekey(delegating_privkey, signer, pub_key_bob, M, N)
 
     capsule_alice1.set_correctness_keys(delegating=delegating_pubkey,
                                         receiving=pub_key_bob,
                                         verifying=signing_privkey.get_pubkey())
 
-    kfrags = pre.split_rekey(delegating_privkey, signer, pub_key_bob, M, N)
-
     cfrags, metadata = [], []
     for i, kfrag in enumerate(kfrags[:M]):
 
@@ -76,34 +76,25 @@ def test_cheating_ursula_replays_old_reencryption(N, M, alices_keys):
         else:
             cfrag = pre.reencrypt(kfrag, capsule_alice1, metadata=metadata_i)
 
-        capsule_alice1.attach_cfrag(cfrag)
+        # Next, we bypass the public method to attach CFrags to the capsule,
+        # -- called Capsule.append(cfrag) -- and insert it directly in the private
+        # list of CFrags. In case you were wondering...DON'T DO THIS!
+        capsule_alice1._attached_cfrags.append(cfrag)
         cfrags.append(cfrag)
 
     # Let's activate the capsule
     capsule_alice1._reconstruct_shamirs_secret(priv_key_bob)
 
     with pytest.raises(pre.GenericUmbralError):
-        sym_key = pre._decapsulate_reencrypted(pub_key_bob.point_key,
-                                               priv_key_bob.bn_key,
-                                               delegating_pubkey.point_key,
-                                               capsule_alice1
-                                               )
-
-    assert not cfrags[0].verify_correctness(capsule_alice1,
-                                            delegating_pubkey,
-                                            signing_privkey.get_pubkey(),
-                                            pub_key_bob,
-                                            )
+        sym_key = pre._decapsulate_reencrypted(priv_key_bob, capsule_alice1)
+
+    assert not cfrags[0].verify_correctness(capsule_alice1)
 
     # The response of cheating Ursula is in cfrags[0],
     # so the rest of CFrags should be correct:
     correct_cases = 0
     for cfrag_i, metadata_i in zip(cfrags[1:], metadata[1:]):
-        if cfrag_i.verify_correctness(capsule_alice1,
-                                      delegating_pubkey,
-                                      signing_privkey.get_pubkey(),
-                                      pub_key_bob,
-                                      ):
+        if cfrag_i.verify_correctness(capsule_alice1):
             correct_cases += 1
         else:
             pytest.fail("One of the cfrags that was supposed to be correct wasn't.")
@@ -112,11 +103,7 @@ def test_cheating_ursula_replays_old_reencryption(N, M, alices_keys):
     # Alternatively, we can try to open the capsule directly.
     # We should get an exception with an attached list of incorrect cfrags
     with pytest.raises(pre.UmbralCorrectnessError) as exception_info:
-        _ = pre._open_capsule(capsule_alice1,
-                              priv_key_bob,
-                              delegating_pubkey,
-                              signing_privkey.get_pubkey(),
-                              )
+        _ = pre._open_capsule(capsule_alice1, priv_key_bob)
     correctness_error = exception_info.value
     assert cfrags[0] in correctness_error.offending_cfrags
     assert len(correctness_error.offending_cfrags) == 1
@@ -133,14 +120,14 @@ def test_cheating_ursula_sends_garbage(N, M, alices_keys):
 
     params = delegating_privkey.params
 
-    sym_key, capsule_alice = pre._encapsulate(delegating_pubkey.point_key, params)
+    sym_key, capsule_alice = pre._encapsulate(delegating_pubkey)
+
+    kfrags = pre.split_rekey(delegating_privkey, signer, pub_key_bob, M, N)
 
     capsule_alice.set_correctness_keys(delegating=delegating_pubkey,
                                        receiving=pub_key_bob,
                                        verifying=signing_privkey.get_pubkey())
 
-    kfrags = pre.split_rekey(delegating_privkey, signer, pub_key_bob, M, N)
-
     cfrags, metadata = [], []
     for i, kfrag in enumerate(kfrags[:M]):
         # Example of potential metadata to describe the re-encryption request
@@ -159,32 +146,20 @@ def test_cheating_ursula_sends_garbage(N, M, alices_keys):
     capsule_alice._reconstruct_shamirs_secret(priv_key_bob)  # activate capsule
 
     with pytest.raises(pre.GenericUmbralError):
-        _unused_key = pre._decapsulate_reencrypted(pub_key_bob.point_key,
-                                                   priv_key_bob.bn_key,
-                                                   delegating_pubkey.point_key,
-                                                   capsule_alice)
+        _unused_key = pre._decapsulate_reencrypted(priv_key_bob, capsule_alice)
 
-    assert not cfrags[0].verify_correctness(capsule_alice,
-                                            delegating_pubkey,
-                                            signing_privkey.get_pubkey(),
-                                            pub_key_bob,
-                                            )
+    assert not cfrags[0].verify_correctness(capsule_alice)
 
     # The response of cheating Ursula is in cfrags[0],
     # so the rest of CFrags chould be correct:
     for cfrag_i, metadata_i in zip(cfrags[1:], metadata[1:]):
-        assert cfrag_i.verify_correctness(capsule_alice,
-                                          delegating_pubkey,
-                                          signing_privkey.get_pubkey(),
-                                          pub_key_bob,
-                                          )
+        assert cfrag_i.verify_correctness(capsule_alice)
 
     # Alternatively, we can try to open the capsule directly.
     # We should get an exception with an attached list of incorrect cfrags
     with pytest.raises(pre.UmbralCorrectnessError) as exception_info:
-        _decapsulated_key = pre._open_capsule(capsule_alice, priv_key_bob, 
-                                              delegating_pubkey,
-                                              signing_privkey.get_pubkey())
+        _decapsulated_key = pre._open_capsule(capsule_alice, priv_key_bob)
+
     correctness_error = exception_info.value
     assert cfrags[0] in correctness_error.offending_cfrags
     assert len(correctness_error.offending_cfrags) == 1
@@ -202,12 +177,12 @@ def test_decryption_fails_when_it_expects_a_proof_and_there_isnt(N, M, alices_ke
     plain_data = b'peace at dawn'
     ciphertext, capsule = pre.encrypt(delegating_privkey.get_pubkey(), plain_data)
 
+    kfrags = pre.split_rekey(delegating_privkey, signer, pub_key_bob, M, N)
+
     capsule.set_correctness_keys(delegating=delegating_privkey.get_pubkey(),
                                  receiving=pub_key_bob,
                                  verifying=signing_privkey.get_pubkey())
-
-    kfrags = pre.split_rekey(delegating_privkey, signer, pub_key_bob, M, N)
-    for kfrag in kfrags:
+    for kfrag in kfrags[:M]:
         cfrag = pre.reencrypt(kfrag, capsule)
         capsule.attach_cfrag(cfrag)
 
@@ -217,9 +192,7 @@ def test_decryption_fails_when_it_expects_a_proof_and_there_isnt(N, M, alices_ke
     cfrag.proof = None
 
     with pytest.raises(cfrag.NoProofProvided):
-        _cleartext = pre.decrypt(ciphertext, capsule, priv_key_bob,
-                                 delegating_privkey.get_pubkey(),
-                                 signing_privkey.get_pubkey())
+        _cleartext = pre.decrypt(ciphertext, capsule, priv_key_bob)
 
 
 @pytest.mark.parametrize("N, M", parameters)
@@ -232,7 +205,7 @@ def test_m_of_n(N, M, alices_keys, bobs_keys):
 
     params = delegating_privkey.params
 
-    sym_key, capsule = pre._encapsulate(delegating_pubkey.point_key, params)
+    sym_key, capsule = pre._encapsulate(delegating_pubkey)
 
     capsule.set_correctness_keys(delegating=delegating_privkey.get_pubkey(),
                                  receiving=pub_key_bob,
@@ -251,10 +224,7 @@ def test_m_of_n(N, M, alices_keys, bobs_keys):
         cfrag = pre.reencrypt(kfrag, capsule, metadata=metadata)
         capsule.attach_cfrag(cfrag)
 
-        assert cfrag.verify_correctness(capsule, delegating_pubkey, 
-                                        signing_privkey.get_pubkey(), pub_key_bob,
-                                        )
+        assert cfrag.verify_correctness(capsule)
 
-    sym_key_from_capsule = pre._open_capsule(capsule, priv_key_bob, delegating_pubkey,
-                                             signing_privkey.get_pubkey())
+    sym_key_from_capsule = pre._open_capsule(capsule, priv_key_bob)
     assert sym_key == sym_key_from_capsule
diff --git a/tests/test_keys/test_key_fragments.py b/tests/test_keys/test_key_fragments.py
index f25a6eab..f35481c4 100644
--- a/tests/test_keys/test_key_fragments.py
+++ b/tests/test_keys/test_key_fragments.py
@@ -32,7 +32,7 @@ def test_cfrag_serialization_with_proof_and_metadata(alices_keys, bobs_keys):
     _receiving_privkey, receiving_pubkey = bobs_keys
 
 
-    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key, delegating_pubkey.params)
+    _unused_key, capsule = pre._encapsulate(delegating_pubkey)
     kfrags = pre.split_rekey(delegating_privkey, signer_alice,
                              receiving_pubkey, 1, 2)
 
@@ -70,7 +70,7 @@ def test_cfrag_serialization_with_proof_but_no_metadata(alices_keys, bobs_keys):
     _receiving_privkey, receiving_pubkey = bobs_keys
     signer_alice = Signer(signing_privkey)
 
-    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key, delegating_pubkey.params)
+    _unused_key, capsule = pre._encapsulate(delegating_pubkey)
     kfrags = pre.split_rekey(delegating_privkey, signer_alice,
                              receiving_pubkey, 1, 2)
 
@@ -108,7 +108,7 @@ def test_cfrag_serialization_no_proof_no_metadata(alices_keys, bobs_keys):
     _receiving_privkey, receiving_pubkey = bobs_keys
     signer_alice = Signer(signing_privkey)
 
-    _unused_key, capsule = pre._encapsulate(delegating_pubkey.point_key, delegating_pubkey.params)
+    _unused_key, capsule = pre._encapsulate(delegating_pubkey)
     kfrags = pre.split_rekey(delegating_privkey, signer_alice,
                              receiving_pubkey, 1, 2)
 
diff --git a/tests/test_simple_api.py b/tests/test_simple_api.py
index 0af43c12..ebc8a94c 100644
--- a/tests/test_simple_api.py
+++ b/tests/test_simple_api.py
@@ -46,8 +46,7 @@ def test_simple_api(N, M, curve=default_curve()):
         cfrag = pre.reencrypt(kfrag, capsule)
         capsule.attach_cfrag(cfrag)
 
-    reenc_cleartext = pre.decrypt(ciphertext, capsule, receiving_privkey,
-                                  delegating_pubkey, signing_pubkey)
+    reenc_cleartext = pre.decrypt(ciphertext, capsule, receiving_privkey)
     assert reenc_cleartext == plain_data
 
 
diff --git a/umbral/_pre.py b/umbral/_pre.py
index cf3ac9a0..138b5ff4 100644
--- a/umbral/_pre.py
+++ b/umbral/_pre.py
@@ -46,10 +46,20 @@ def prove_cfrag_correctness(cfrag: "CapsuleFrag",
 
 
 def assess_cfrag_correctness(cfrag,
-                             capsule: "Capsule",
-                             delegating_point,
-                             signing_pubkey,
-                             receiving_point):
+                             capsule: "Capsule"):
+
+    correctness_keys = capsule.get_correctness_keys()
+
+    delegating_pubkey = correctness_keys['delegating']
+    signing_pubkey = correctness_keys['verifying']
+    receiving_pubkey = correctness_keys['receiving']
+
+    if not all((delegating_pubkey, signing_pubkey, receiving_pubkey)):
+        raise TypeError("Need all three keys to verify correctness.")
+
+    delegating_point = delegating_pubkey.point_key
+    receiving_point = receiving_pubkey.point_key
+
     params = capsule._umbral_params
 
     ####
diff --git a/umbral/fragments.py b/umbral/fragments.py
index b9d8f259..2afc4619 100644
--- a/umbral/fragments.py
+++ b/umbral/fragments.py
@@ -220,19 +220,8 @@ def to_bytes(self):
 
         return serialized_cfrag
 
-    def verify_correctness(self,
-                           capsule: "Capsule",
-                           delegating_pubkey: UmbralPublicKey,
-                           signing_pubkey,
-                           receiving_pubkey: UmbralPublicKey):
-        if not all((delegating_pubkey, signing_pubkey, receiving_pubkey)):
-            raise TypeError("Need all three keys to verify correctness.")
-
-        pubkey_a_point = delegating_pubkey.point_key
-        pubkey_b_point = receiving_pubkey.point_key
-
-        return assess_cfrag_correctness(self, capsule, pubkey_a_point,
-                                        signing_pubkey, pubkey_b_point)
+    def verify_correctness(self, capsule: "Capsule"):
+        return assess_cfrag_correctness(self, capsule)
 
     def attach_proof(self, e2, v2, u1, u2, z3, kfrag_signature, metadata):
         self.proof = CorrectnessProof(point_e2=e2,
diff --git a/umbral/pre.py b/umbral/pre.py
index 82ef4aeb..3701adbe 100644
--- a/umbral/pre.py
+++ b/umbral/pre.py
@@ -139,6 +139,10 @@ def _set_cfrag_correctness_key(self, key_type, key: UmbralPublicKey):
         else:
             raise ValueError("The {} key is already set; you can't set it again.".format(key_type))
 
+
+    def get_correctness_keys(self):
+        return dict(self._cfrag_correctness_keys)
+
     def set_correctness_keys(self,
                  delegating: UmbralPublicKey = None,
                  receiving: UmbralPublicKey = None,
@@ -189,17 +193,12 @@ def original_components(self) -> Tuple[Point, Point, CurveBN]:
     def activated_components(self) -> Union[Tuple[None, None, None], Tuple[Point, Point, Point]]:
         return self._point_e_prime, self._point_v_prime, self._point_noninteractive
 
-    def _reconstruct_shamirs_secret(self,
-                                    priv_b: Union[UmbralPrivateKey, CurveBN]
-                                    ) -> None:
+    def _reconstruct_shamirs_secret(self, priv_b: UmbralPrivateKey) -> None:
         params = self._umbral_params
         g = params.g
 
-        if isinstance(priv_b, UmbralPrivateKey):
-            pub_b = priv_b.get_pubkey()
-            priv_b = priv_b.bn_key
-        else:
-            pub_b = priv_b * g
+        pub_b = priv_b.get_pubkey()
+        priv_b = priv_b.bn_key
 
         cfrag_0 = self._attached_cfrags[0]
         id_0 = cfrag_0._kfrag_id
@@ -369,10 +368,11 @@ def reencrypt(kfrag: KFrag, capsule: Capsule, provide_proof = True,
     return cfrag
 
 
-def _encapsulate(alice_pubkey: Point, params: UmbralParameters,
+def _encapsulate(alice_pubkey: UmbralPublicKey, 
                  key_length = DEM_KEYSIZE) -> Tuple[bytes, Capsule]:
     """Generates a symmetric key and its associated KEM ciphertext"""
 
+    params = alice_pubkey.params
     g = params.g
 
     priv_r = CurveBN.gen_rand(params.curve)
@@ -384,7 +384,7 @@ def _encapsulate(alice_pubkey: Point, params: UmbralParameters,
     h = CurveBN.hash(pub_r, pub_u, params=params)
     s = priv_u + (priv_r * h)
 
-    shared_key = (priv_r + priv_u) * alice_pubkey
+    shared_key = (priv_r + priv_u) * alice_pubkey.point_key
 
     # Key to be used for symmetric encryption
     key = kdf(shared_key, key_length)
@@ -392,10 +392,12 @@ def _encapsulate(alice_pubkey: Point, params: UmbralParameters,
     return key, Capsule(point_e=pub_r, point_v=pub_u, bn_sig=s, params=params)
 
 
-def _decapsulate_original(priv_key: CurveBN, capsule: Capsule, 
+def _decapsulate_original(priv_key: UmbralPrivateKey, capsule: Capsule, 
                           key_length = DEM_KEYSIZE) -> bytes:
     """Derive the same symmetric key"""
 
+    priv_key = priv_key.bn_key
+
     shared_key = priv_key * (capsule._point_e + capsule._point_v)
     key = kdf(shared_key, key_length)
 
@@ -407,12 +409,14 @@ def _decapsulate_original(priv_key: CurveBN, capsule: Capsule,
     return key
 
 
-def _decapsulate_reencrypted(pub_key: Point, priv_key: CurveBN,
-                             orig_pub_key: Point, capsule: Capsule,
+def _decapsulate_reencrypted(receiving_privkey: UmbralPrivateKey, capsule: Capsule,
                              key_length = DEM_KEYSIZE) -> bytes:
     """Derive the same symmetric key"""
     params = capsule._umbral_params
 
+    pub_key = receiving_privkey.get_pubkey().point_key
+    priv_key = receiving_privkey.bn_key
+
     ni = capsule._point_noninteractive
     d = CurveBN.hash(ni, pub_key, priv_key * ni, params=params)
 
@@ -428,6 +432,7 @@ def _decapsulate_reencrypted(pub_key: Point, priv_key: CurveBN,
     s = capsule._bn_sig
     h = CurveBN.hash(e, v, params=params)
     inv_d = ~d
+    orig_pub_key = capsule.get_correctness_keys()['delegating'].point_key
 
     if not (s * inv_d) * orig_pub_key == (h * e_prime) + v_prime:
         raise GenericUmbralError()
@@ -441,9 +446,7 @@ def encrypt(alice_pubkey: UmbralPublicKey, plaintext: bytes) -> Tuple[bytes, Cap
 
     Returns the ciphertext and the KEM Capsule.
     """
-    params = alice_pubkey.params
-
-    key, capsule = _encapsulate(alice_pubkey.point_key, params, DEM_KEYSIZE)
+    key, capsule = _encapsulate(alice_pubkey, DEM_KEYSIZE)
 
     capsule_bytes = bytes(capsule)
 
@@ -455,8 +458,6 @@ def encrypt(alice_pubkey: UmbralPublicKey, plaintext: bytes) -> Tuple[bytes, Cap
 
 def _open_capsule(capsule: Capsule,
                   receiving_privkey: UmbralPrivateKey,
-                  delegating_pubkey: UmbralPublicKey,
-                  signing_pubkey: UmbralPublicKey,
                   check_proof=True) -> bytes:
     """
     Activates the Capsule from the attached CFrags,
@@ -466,33 +467,26 @@ def _open_capsule(capsule: Capsule,
     """
 
     receiving_pubkey = receiving_privkey.get_pubkey()
-    priv_b = receiving_privkey.bn_key
 
     if check_proof:
         offending_cfrags = []
         for cfrag in capsule._attached_cfrags:
-            if not cfrag.verify_correctness(capsule=capsule,
-                                            delegating_pubkey=delegating_pubkey,
-                                            signing_pubkey=signing_pubkey,
-                                            receiving_pubkey=receiving_pubkey):
+            if not cfrag.verify_correctness(capsule):
                 offending_cfrags.append(cfrag)
 
         if offending_cfrags:
             error_msg = "Decryption error: Some CFrags are not correct"
             raise UmbralCorrectnessError(error_msg, offending_cfrags)
 
-    capsule._reconstruct_shamirs_secret(priv_b)
+    capsule._reconstruct_shamirs_secret(receiving_privkey)
 
-    key = _decapsulate_reencrypted(receiving_pubkey.point_key, priv_b, 
-                                   delegating_pubkey.point_key, capsule)
+    key = _decapsulate_reencrypted(receiving_privkey, capsule)
     return key
 
 
 def decrypt(ciphertext: bytes,
             capsule: Capsule,
             decrypting_key: UmbralPrivateKey,
-            delegating_pubkey: UmbralPublicKey = None,
-            verifying_key = None,
             check_proof=True) -> bytes:
     """
     Opens the capsule and gets what's inside.
@@ -505,9 +499,7 @@ def decrypt(ciphertext: bytes,
         # Since there are cfrags attached, we assume this is Bob opening the Capsule.
         # (i.e., this is a re-encrypted capsule)
 
-        encapsulated_key = _open_capsule(capsule, decrypting_key, 
-                                         delegating_pubkey, verifying_key,
-                                         check_proof=check_proof)
+        encapsulated_key = _open_capsule(capsule, decrypting_key, check_proof=check_proof)
         dem = UmbralDEM(encapsulated_key)
 
         original_capsule_bytes = capsule._original_to_bytes()
@@ -515,7 +507,7 @@ def decrypt(ciphertext: bytes,
     else:
         # Since there aren't cfrags attached, we assume this is Alice opening the Capsule.
         # (i.e., this is an original capsule)
-        decapsulated_key = _decapsulate_original(decrypting_key.bn_key, capsule)
+        decapsulated_key = _decapsulate_original(decrypting_key, capsule)
         dem = UmbralDEM(decapsulated_key)
 
         capsule_bytes = bytes(capsule)

From 9e126bde293224a3e1b14439a962e91c27a52db0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Mon, 4 Jun 2018 14:38:00 +0200
Subject: [PATCH 08/11] Fixes bug in Capsule.attach_cfrag that allowed to
 attach incorrect CFrag

---
 umbral/pre.py | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/umbral/pre.py b/umbral/pre.py
index 3701adbe..e3db8bee 100644
--- a/umbral/pre.py
+++ b/umbral/pre.py
@@ -177,15 +177,11 @@ def verify(self) -> bool:
         return s * g == v + (h * e)
 
     def attach_cfrag(self, cfrag: CapsuleFrag) -> None:
-        self.verify_cfrag(cfrag)
-        self._attached_cfrags.append(cfrag)
-
-    def verify_cfrag(self, cfrag):
-        return cfrag.verify_correctness(self,
-                                        self._cfrag_correctness_keys["delegating"],
-                                        self._cfrag_correctness_keys["receiving"],
-                                        self._cfrag_correctness_keys["verifying"],
-                                        )
+        if cfrag.verify_correctness(self):
+            self._attached_cfrags.append(cfrag)
+        else:
+            error_msg = "CFrag is not correct and cannot be attached to the Capsule"
+            raise UmbralCorrectnessError(error_msg, [cfrag])
 
     def original_components(self) -> Tuple[Point, Point, CurveBN]:
         return self._point_e, self._point_v, self._bn_sig

From d70d21dfe0e111c368542bde0e05439605ec3b81 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Tue, 5 Jun 2018 03:49:22 +0200
Subject: [PATCH 09/11] Fix some functions definition according to PEP8

---
 umbral/_pre.py |  3 +--
 umbral/pre.py  | 17 +++++++----------
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/umbral/_pre.py b/umbral/_pre.py
index 138b5ff4..f2b2cc08 100644
--- a/umbral/_pre.py
+++ b/umbral/_pre.py
@@ -45,8 +45,7 @@ def prove_cfrag_correctness(cfrag: "CapsuleFrag",
         raise capsule.NotValid("Capsule verification failed.")
 
 
-def assess_cfrag_correctness(cfrag,
-                             capsule: "Capsule"):
+def assess_cfrag_correctness(cfrag, capsule: "Capsule"):
 
     correctness_keys = capsule.get_correctness_keys()
 
diff --git a/umbral/pre.py b/umbral/pre.py
index e3db8bee..00f520b8 100644
--- a/umbral/pre.py
+++ b/umbral/pre.py
@@ -39,7 +39,7 @@ def __init__(self,
                  point_noninteractive=None,
                  delegating_pubkey: UmbralPublicKey = None,
                  receiving_pubkey: UmbralPublicKey = None,
-                 verifying_pubkey = None):
+                 verifying_pubkey=None):
 
         self._umbral_params = params
 
@@ -146,7 +146,7 @@ def get_correctness_keys(self):
     def set_correctness_keys(self,
                  delegating: UmbralPublicKey = None,
                  receiving: UmbralPublicKey = None,
-                 verifying: UmbralPublicKey = None
+                 verifying=None
                  ):
         delegating_key_details = self._set_cfrag_correctness_key("delegating", delegating)
         receiving_key_details = self._set_cfrag_correctness_key("receiving", receiving)
@@ -365,7 +365,7 @@ def reencrypt(kfrag: KFrag, capsule: Capsule, provide_proof = True,
 
 
 def _encapsulate(alice_pubkey: UmbralPublicKey, 
-                 key_length = DEM_KEYSIZE) -> Tuple[bytes, Capsule]:
+                 key_length=DEM_KEYSIZE) -> Tuple[bytes, Capsule]:
     """Generates a symmetric key and its associated KEM ciphertext"""
 
     params = alice_pubkey.params
@@ -389,7 +389,7 @@ def _encapsulate(alice_pubkey: UmbralPublicKey,
 
 
 def _decapsulate_original(priv_key: UmbralPrivateKey, capsule: Capsule, 
-                          key_length = DEM_KEYSIZE) -> bytes:
+                          key_length=DEM_KEYSIZE) -> bytes:
     """Derive the same symmetric key"""
 
     priv_key = priv_key.bn_key
@@ -406,7 +406,7 @@ def _decapsulate_original(priv_key: UmbralPrivateKey, capsule: Capsule,
 
 
 def _decapsulate_reencrypted(receiving_privkey: UmbralPrivateKey, capsule: Capsule,
-                             key_length = DEM_KEYSIZE) -> bytes:
+                             key_length=DEM_KEYSIZE) -> bytes:
     """Derive the same symmetric key"""
     params = capsule._umbral_params
 
@@ -452,8 +452,7 @@ def encrypt(alice_pubkey: UmbralPublicKey, plaintext: bytes) -> Tuple[bytes, Cap
     return ciphertext, capsule
 
 
-def _open_capsule(capsule: Capsule,
-                  receiving_privkey: UmbralPrivateKey,
+def _open_capsule(capsule: Capsule, receiving_privkey: UmbralPrivateKey,
                   check_proof=True) -> bytes:
     """
     Activates the Capsule from the attached CFrags,
@@ -480,9 +479,7 @@ def _open_capsule(capsule: Capsule,
     return key
 
 
-def decrypt(ciphertext: bytes,
-            capsule: Capsule,
-            decrypting_key: UmbralPrivateKey,
+def decrypt(ciphertext: bytes, capsule: Capsule, decrypting_key: UmbralPrivateKey,
             check_proof=True) -> bytes:
     """
     Opens the capsule and gets what's inside.

From d34fd6e3e7b7316a5cd5df48d4927677674a0935 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Tue, 5 Jun 2018 03:49:53 +0200
Subject: [PATCH 10/11] Fix __bytes__ definition in CorrectnessProof

---
 umbral/fragments.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/umbral/fragments.py b/umbral/fragments.py
index 2afc4619..2fc36672 100644
--- a/umbral/fragments.py
+++ b/umbral/fragments.py
@@ -149,7 +149,7 @@ def to_bytes(self) -> bytes:
 
         return result
 
-    def _bytes__(self):
+    def __bytes__(self):
         return self.to_bytes()
 
 

From fa4375d9dc0c31bfcf48f727cdce702e7c6f43db Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Tue, 5 Jun 2018 03:51:30 +0200
Subject: [PATCH 11/11] Adapt verify_kfrag to simplified API (I missed this in
 the main commit)

---
 umbral/_pre.py      | 13 ++++++++++---
 umbral/fragments.py |  8 +-------
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/umbral/_pre.py b/umbral/_pre.py
index f2b2cc08..de9de8e7 100644
--- a/umbral/_pre.py
+++ b/umbral/_pre.py
@@ -111,14 +111,21 @@ def assess_cfrag_correctness(cfrag, capsule: "Capsule"):
 
 
 def verify_kfrag(kfrag,
-                 delegating_point,
+                 delegating_pubkey: UmbralPublicKey,
                  signing_pubkey,
-                 receiving_point,
-                 params: UmbralParameters
+                 receiving_pubkey: UmbralPublicKey
                  ):
 
+
+    params = delegating_pubkey.params
+    if not params == receiving_pubkey.params:
+        raise ValueError("The delegating and receiving keys must use the same UmbralParameters")
+
     u = params.u
 
+    delegating_point = delegating_pubkey.point_key
+    receiving_point = receiving_pubkey.point_key
+
     id = kfrag._id
     key = kfrag._bn_key
     u1 = kfrag._point_commitment
diff --git a/umbral/fragments.py b/umbral/fragments.py
index 2fc36672..9034e21e 100644
--- a/umbral/fragments.py
+++ b/umbral/fragments.py
@@ -72,13 +72,7 @@ def verify(self,
                delegating_pubkey: UmbralPublicKey,
                receiving_pubkey: UmbralPublicKey):
 
-        params = delegating_pubkey.params
-
-        return verify_kfrag(self,
-                            delegating_pubkey.point_key,
-                            signing_pubkey,
-                            receiving_pubkey.point_key,
-                            params)
+        return verify_kfrag(self, delegating_pubkey, signing_pubkey, receiving_pubkey)
 
     def __bytes__(self):
         return self.to_bytes()