From 32c5f542e67b74d12c71198218f7036c89db382b Mon Sep 17 00:00:00 2001
From: tuxxy <john@nucypher.com>
Date: Mon, 23 Apr 2018 12:54:06 -0600
Subject: [PATCH] Add get_size classmethods for CurveBN, Point, fragments, and
 CorrectnessProof.

Fix syntax error
---
 umbral/curvebn.py   |  9 +++++++++
 umbral/fragments.py | 36 ++++++++++++++++++++++++++++++++++++
 umbral/point.py     |  9 +++++++++
 umbral/pre.py       |  2 +-
 umbral/utils.py     |  2 ++
 5 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/umbral/curvebn.py b/umbral/curvebn.py
index 4a07502..acdf375 100644
--- a/umbral/curvebn.py
+++ b/umbral/curvebn.py
@@ -20,6 +20,15 @@ class CurveBN(object):
         self.group = group
         self.order = order
 
+    @classmethod
+    def get_size(cls, curve: ec.EllipticCurve=None):
+        """
+        Returns the size (in bytes) of a CurveBN given the curve.
+        If no curve is provided, it uses the default.
+        """
+        curve = curve if curve is not None else default_curve()
+        return get_curve_keysize_bytes(curve)
+
     @classmethod
     def gen_rand(cls, curve: ec.EllipticCurve = None):
         """
diff --git a/umbral/fragments.py b/umbral/fragments.py
index 16f282c..d8016c9 100644
--- a/umbral/fragments.py
+++ b/umbral/fragments.py
@@ -18,6 +18,18 @@ class KFrag(object):
         self._bn_sig1 = bn_sig1
         self._bn_sig2 = bn_sig2
 
+    @classmethod
+    def get_size(cls, curve: ec.EllipticCurve=None):
+        """
+        Returns the size (in bytes) of a KFrag given the curve.
+        If no curve is provided, it will use the default curve.
+        """
+        curve = curve if curve is not None else default_curve()
+        bn_size = CurveBN.get_size(curve)
+        point_size = Point.get_size(curve)
+
+        return (bn_size * 4) + (point_size * 2)
+
     @classmethod
     def from_bytes(cls, data: bytes, curve: ec.EllipticCurve = None):
         """
@@ -90,6 +102,18 @@ class CorrectnessProof(object):
         self._bn_sig = bn_sig
         self.metadata = metadata
 
+    @classmethod
+    def get_size(cls, curve: ec.EllipticCurve=None):
+        """
+        Returns the size (in bytes) of a CorrectnessProof without the metadata.
+        If no curve is given, it will use the default curve.
+        """
+        curve = curve if curve is not None else default_curve()
+        bn_size = CurveBN.get_size(curve=curve)
+        point_size = Point.get_size(curve=curve)
+
+        return (bn_size * 3) + (point_size * 4)
+
     @classmethod
     def from_bytes(cls, data: bytes, curve: ec.EllipticCurve=None):
         """
@@ -151,6 +175,18 @@ class CapsuleFrag(object):
         self._point_noninteractive = point_noninteractive
         self.proof = proof
 
+    @classmethod
+    def get_size(cls, curve: ec.EllipticCurve=None):
+        """
+        Returns the size (in bytes) of a CapsuleFrag given the curve.
+        If no curve is provided, it will use the default curve.
+        """
+        curve = curve if curve is not None else default_curve()
+        bn_size = CurveBN.get_size(curve)
+        point_size = Point.get_size(curve)
+
+        return (bn_size * 1) + (point_size * 3)
+
     @classmethod
     def from_bytes(cls, data: bytes, curve: ec.EllipticCurve = None):
         """
diff --git a/umbral/point.py b/umbral/point.py
index 0562773..b446f02 100644
--- a/umbral/point.py
+++ b/umbral/point.py
@@ -19,6 +19,15 @@ class Point(object):
         self.curve_nid = curve_nid
         self.group = group
 
+    @classmethod
+    def get_size(cls, curve: ec.EllipticCurve=None):
+        """
+        Returns the size (in bytes) of a compressed Point given a curve.
+        If no curve is provided, it uses the default curve.
+        """
+        curve = curve if curve is not None else default_curve()
+        return get_curve_keysize_bytes(curve)
+
     @classmethod
     def gen_rand(cls, curve: ec.EllipticCurve=None):
         """
diff --git a/umbral/pre.py b/umbral/pre.py
index 65fa58e..fa48c1f 100644
--- a/umbral/pre.py
+++ b/umbral/pre.py
@@ -123,7 +123,7 @@ class Capsule(object):
     def original_components(self) -> Tuple[Point, Point, CurveBN]:
         return self._point_e, self._point_v, self._bn_sig
 
-    def activated_components(self) -> Union[Tuple[None, NCurveBNone], Tuple[Point, Point, Point]]:
+    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, 
diff --git a/umbral/utils.py b/umbral/utils.py
index 32f8939..71e10ec 100644
--- a/umbral/utils.py
+++ b/umbral/utils.py
@@ -41,4 +41,6 @@ def kdf(ecpoint, key_length):
 
 
 def get_curve_keysize_bytes(curve):
+    # We use the ceil operation to fit all bytes on curve sizes where eight is
+    # not evenly divisible.
     return int(math.ceil(curve.key_size / 8.00))