format using black

.github/workflows/ci.yml

@@ -7,6 +7,38 @@ on:
   pull_request:
 
 jobs:
+  code-format:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 10
+      - name: Verify git status
+        run: |
+          git status
+          git remote -v
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Display Python version
+        run: python -c "import sys; print(sys.version)"
+      - name: Display installed python package versions
+        run: |
+          pip list || :
+      - name: Install build dependencies
+        run: |
+          pip install -r requirements.txt
+      - name: Install black
+        run: |
+          pip install black==24.4.2
+      - name: Display installed python package versions
+        run: |
+          pip list || :
+      - name: Check formatting
+        run: |
+          black --check --line-length 79 .
+
   test:
     runs-on: ${{ matrix.os }}
     strategy:

aes256_ctr_drbg.py

@@ -2,18 +2,19 @@
 from utils import xor_bytes
 from Crypto.Cipher import AES
 
+
 class AES256_CTR_DRBG:
     def __init__(self, seed=None, personalization=b""):
         self.seed_length = 48
         self.reseed_interval = 2**48
-        self.key = bytes([0])*32
-        self.V   = bytes([0])*16
+        self.key = bytes([0]) * 32
+        self.V = bytes([0]) * 16
         self.entropy_input = self.__check_entropy_input(seed)
-      
+
         seed_material = self.__instantiate(personalization=personalization)
         self.ctr_drbg_update(seed_material)
         self.reseed_ctr = 1
-        
+
     def __check_entropy_input(self, entropy_input):
         """
         If no entropy given, us os.urandom, else
@@ -22,76 +23,86 @@ def __check_entropy_input(self, entropy_input):
         if entropy_input is None:
             return os.urandom(self.seed_length)
         elif len(entropy_input) != self.seed_length:
-            raise ValueError(f"The entropy input must be of length: {self.seed_length}. Input has length {len(entropy_input)}")
+            raise ValueError(
+                f"The entropy input must be of length: {self.seed_length}. "
+                f"Input has length {len(entropy_input)}"
+            )
         return entropy_input
-        
+
     def __instantiate(self, personalization=b""):
         """
         Combine the input seed and optional personalisation
         string into the seed material for the DRBG
         """
         if len(personalization) > self.seed_length:
-            raise ValueError(f"The Personalization String must be at most length: {self.seed_length}. Input has length {len(personalization)}")
+            raise ValueError(
+                f"The Personalization String must be at most length: "
+                f"{self.seed_length}. Input has length {len(personalization)}"
+            )
         elif len(personalization) < self.seed_length:
-            personalization += bytes([0]) * (self.seed_length - len(personalization))
+            personalization += bytes([0]) * (
+                self.seed_length - len(personalization)
+            )
         # debugging
         assert len(personalization) == self.seed_length
         return xor_bytes(self.entropy_input, personalization)
 
     def __increment_counter(self):
-        int_V = int.from_bytes(self.V, 'big')
-        new_V = (int_V + 1) % 2**(8*16)
-        self.V = new_V.to_bytes(16, byteorder='big')
-        
+        int_V = int.from_bytes(self.V, "big")
+        new_V = (int_V + 1) % 2 ** (8 * 16)
+        self.V = new_V.to_bytes(16, byteorder="big")
+
     def ctr_drbg_update(self, provided_data):
         tmp = b""
         cipher = AES.new(self.key, AES.MODE_ECB)
         # Collect bytes from AES ECB
         while len(tmp) != self.seed_length:
             self.__increment_counter()
-            tmp  += cipher.encrypt(self.V)
-        
+            tmp += cipher.encrypt(self.V)
+
         # Take the first 48 bytes
-        tmp = tmp[:self.seed_length]
+        tmp = tmp[: self.seed_length]
         tmp = xor_bytes(tmp, provided_data)
-        
+
         # Set the new values of key and V
         self.key = tmp[:32]
         self.V = tmp[32:]
-        
+
     def reseed(self, additional_information=b""):
         """
-        Reseed the DRBG for when reseed_ctr hits the 
+        Reseed the DRBG for when reseed_ctr hits the
         limit.
         """
         seed_material = self.__instantiate(additional_information)
         self.ctr_drbg_update(seed_material)
         self.reseed_ctr = 1
-        
+
     def random_bytes(self, num_bytes, additional=None):
         if self.reseed_ctr >= self.reseed_interval:
             raise Warning("The DRBG has been exhausted! Reseed!")
-        
+
         # Set the optional additional information
         if additional is None:
             additional = bytes([0]) * self.seed_length
         else:
             if len(additional) > self.seed_length:
-                 raise ValueError(f"The additional input must be of length at most: {self.seed_length}. Input has length {len(seed)}")
+                raise ValueError(
+                    f"The additional input must be of length at most: "
+                    f"{self.seed_length}. Input has length {len(seed)}"
+                )
             elif len(additional) < self.seed_length:
                 additional += bytes([0]) * (self.seed_length - len(additional))
             self.ctr_drbg_update(additional)
-        
+
         # Collect bytes!
         tmp = b""
         cipher = AES.new(self.key, AES.MODE_ECB)
         while len(tmp) < num_bytes:
             self.__increment_counter()
             tmp += cipher.encrypt(self.V)
-        
+
         # Collect only the requested number of bits
         output_bytes = tmp[:num_bytes]
         self.ctr_drbg_update(additional)
         self.reseed_ctr += 1
         return output_bytes
-

baby-kyber.py

@@ -10,85 +10,89 @@
 from polynomials import *
 from modules import *
 
+
 def keygen():
     # Randomness fixed for example
     # Generate a secret key which
-    # is a vector with elements 
+    # is a vector with elements
     # pulled from a centred
     # binomial distribution
-    s0 = R([0,1,-1,-1])
-    s1 = R([0,-1,0,-1])
-    s = M([s0,s1]).transpose()
+    s0 = R([0, 1, -1, -1])
+    s1 = R([0, -1, 0, -1])
+    s = M([s0, s1]).transpose()
 
     # Randomness fixed for example
-    # Generate a 2x2 matrix with 
+    # Generate a 2x2 matrix with
     # elements taken randomly from
     # R_q
-    A00 = R([11,16,16,6])
-    A01 = R([3,6,4,9])
-    A10 = R([1,10,3,5])
-    A11 = R([15,9,1,6])
-    A = M([[A00, A01],[A10, A11]])
+    A00 = R([11, 16, 16, 6])
+    A01 = R([3, 6, 4, 9])
+    A10 = R([1, 10, 3, 5])
+    A11 = R([15, 9, 1, 6])
+    A = M([[A00, A01], [A10, A11]])
 
     # Randomness fixed for example
     # gen random vector `e` from
     # binomial distribution
-    e0 = R([0,0,1,0])
-    e1 = R([0,-1,1,0])
-    e = M([e0,e1]).transpose()
+    e0 = R([0, 0, 1, 0])
+    e1 = R([0, -1, 1, 0])
+    e = M([e0, e1]).transpose()
 
     # Compute `t` from example
     t = A @ s + e
 
     # Check against blogpost
-    assert t == M([R([7,0,15,16]),R([6,11,12,10])]).transpose()
+    assert t == M([R([7, 0, 15, 16]), R([6, 11, 12, 10])]).transpose()
     return (A, t), s
 
+
 def enc(m, public_key):
     # randomness fixed for example
     # gen random vector `r` from
     # binomial distribution
-    r0 = R([0,0,1,-1]) 
-    r1 = R([-1,0,1,1])
+    r0 = R([0, 0, 1, -1])
+    r1 = R([-1, 0, 1, 1])
     r = M([r0, r1]).transpose()
-    
+
     # randomness fixed for example
     # gen random vector `e_1` from
     # binomial distribution
-    e_10 = R([0,1,1,0])
-    e_11 = R([0,0,1,0])
+    e_10 = R([0, 1, 1, 0])
+    e_11 = R([0, 0, 1, 0])
     e_1 = M([e_10, e_11]).transpose()
-    
+
     # randomness fixed for example
     # gen random polynomial `e_2` from
     # binomial distribution
-    e_2 = R([0,0,-1,-1])
+    e_2 = R([0, 0, -1, -1])
 
     A, t = public_key
     poly_m = R.decode(m).decompress(1)
     # Check against blogpost
-    assert poly_m == R([9,9,0,9])
-    
+    assert poly_m == R([9, 9, 0, 9])
+
     u = A.transpose() @ r + e_1
     # Check against blogpost
-    assert u == M([R([3,10,11,11]), R([11,13,4,4])]).transpose()
-    
+    assert u == M([R([3, 10, 11, 11]), R([11, 13, 4, 4])]).transpose()
+
     # Typo in blog post, we need to use
     # `- m` rather than `+ m` to make values match
-    v = t.dot(r) + e_2 - poly_m  
-    assert v == R([15, 8 , 6, 7])
+    v = t.dot(r) + e_2 - poly_m
+    assert v == R([15, 8, 6, 7])
     return u, v
 
+
 def dec(u, v, s):
     m_n = v - s.dot(u)
     # Check against blogpost
-    assert m_n == R([5,7,14,7])
+    assert m_n == R([5, 7, 14, 7])
     # Check against blogpost
     m_n_reduced = m_n.compress(1)
-    assert m_n_reduced == R([1,1,0,1])
+    assert m_n_reduced == R([1, 1, 0, 1])
     return m_n_reduced.encode(l=2)
-
-if __name__ == '__main__':
+
+
+if __name__ == "__main__":
     R = PolynomialRing(17, 4)
     M = Module(R)
     # Our polynomial encoding follows Kyber spec
@@ -97,13 +101,11 @@ def dec(u, v, s):
     # >>> R.decode(bytes([69]))
     # 1 + x + x^3
     m = bytes([69])
-    assert R.decode(m) == R([1,1,0,1])
+    assert R.decode(m) == R([1, 1, 0, 1])
     # Generate keypair
     pub, priv = keygen()
     # Encrypt message
     u, v = enc(m, pub)
     # Decrypt message
     n = dec(u, v, priv)
     assert n == m
-
-