Added type annotations + some fixes to get them correct

One functional change: `CryptoOperation.read_infile()` now reads bytes
from `sys.stdin` instead of text. This is necessary to be consistent with
the rest of the code, which all deals with bytes.

Author: sybrenstuvel

rsa/_compat.py

@@ -21,14 +21,11 @@
 from struct import pack
 
 
-def byte(num):## XXX
+def byte(num: int):
     """
     Converts a number between 0 and 255 (both inclusive) to a base-256 (byte)
     representation.
 
-    Use it as a replacement for ``chr`` where you are expecting a byte
-    because this will work on all current versions of Python::
-
     :param num:
         An unsigned integer between 0 and 255 (both inclusive).
     :returns:
@@ -37,7 +34,7 @@ def byte(num):## XXX
     return pack("B", num)
 
 
-def xor_bytes(b1, b2):
+def xor_bytes(b1: bytes, b2: bytes) -> bytes:
     """
     Returns the bitwise XOR result between two bytes objects, b1 ^ b2.
 

rsa/cli.py

@@ -22,20 +22,21 @@
 import abc
 import sys
 import typing
-from optparse import OptionParser
+import optparse
 
 import rsa
 import rsa.key
 import rsa.pkcs1
 
 HASH_METHODS = sorted(rsa.pkcs1.HASH_METHODS.keys())
+Indexable = typing.Union[typing.Tuple, typing.List[str]]
 
 
-def keygen():
+def keygen() -> None:
     """Key generator."""
 
     # Parse the CLI options
-    parser = OptionParser(usage='usage: %prog [options] keysize',
+    parser = optparse.OptionParser(usage='usage: %prog [options] keysize',
                           description='Generates a new RSA keypair of "keysize" bits.')
 
     parser.add_option('--pubout', type='string',
@@ -104,21 +105,22 @@ class CryptoOperation(metaclass=abc.ABCMeta):
 
     key_class = rsa.PublicKey  # type: typing.Type[rsa.key.AbstractKey]
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.usage = self.usage % self.__class__.__dict__
         self.input_help = self.input_help % self.__class__.__dict__
         self.output_help = self.output_help % self.__class__.__dict__
 
     @abc.abstractmethod
-    def perform_operation(self, indata, key, cli_args):
+    def perform_operation(self, indata: bytes, key: rsa.key.AbstractKey,
+                          cli_args: Indexable):
         """Performs the program's operation.
 
         Implement in a subclass.
 
         :returns: the data to write to the output.
         """
 
-    def __call__(self):
+    def __call__(self) -> None:
         """Runs the program."""
 
         (cli, cli_args) = self.parse_cli()
@@ -133,13 +135,13 @@ def __call__(self):
         if self.has_output:
             self.write_outfile(outdata, cli.output)
 
-    def parse_cli(self):
+    def parse_cli(self) -> typing.Tuple[optparse.Values, typing.List[str]]:
         """Parse the CLI options
 
         :returns: (cli_opts, cli_args)
         """
 
-        parser = OptionParser(usage=self.usage, description=self.description)
+        parser = optparse.OptionParser(usage=self.usage, description=self.description)
 
         parser.add_option('-i', '--input', type='string', help=self.input_help)
 
@@ -158,7 +160,7 @@ def parse_cli(self):
 
         return cli, cli_args
 
-    def read_key(self, filename, keyform):
+    def read_key(self, filename: str, keyform: str) -> rsa.key.AbstractKey:
         """Reads a public or private key."""
 
         print('Reading %s key from %s' % (self.keyname, filename), file=sys.stderr)
@@ -167,7 +169,7 @@ def read_key(self, filename, keyform):
 
         return self.key_class.load_pkcs1(keydata, keyform)
 
-    def read_infile(self, inname):
+    def read_infile(self, inname: str) -> bytes:
         """Read the input file"""
 
         if inname:
@@ -176,9 +178,9 @@ def read_infile(self, inname):
                 return infile.read()
 
         print('Reading input from stdin', file=sys.stderr)
-        return sys.stdin.read()
+        return sys.stdin.buffer.read()
 
-    def write_outfile(self, outdata, outname):
+    def write_outfile(self, outdata: bytes, outname: str) -> None:
         """Write the output file"""
 
         if outname:
@@ -200,9 +202,10 @@ class EncryptOperation(CryptoOperation):
     operation_past = 'encrypted'
     operation_progressive = 'encrypting'
 
-    def perform_operation(self, indata, pub_key, cli_args=None):
+    def perform_operation(self, indata: bytes, pub_key: rsa.key.AbstractKey,
+            cli_args: Indexable=()):
         """Encrypts files."""
-
+        assert isinstance(pub_key, rsa.key.PublicKey)
         return rsa.encrypt(indata, pub_key)
 
 
@@ -217,9 +220,10 @@ class DecryptOperation(CryptoOperation):
     operation_progressive = 'decrypting'
     key_class = rsa.PrivateKey
 
-    def perform_operation(self, indata, priv_key, cli_args=None):
+    def perform_operation(self, indata: bytes, priv_key: rsa.key.AbstractKey,
+            cli_args: Indexable=()):
         """Decrypts files."""
-
+        assert isinstance(priv_key, rsa.key.PrivateKey)
         return rsa.decrypt(indata, priv_key)
 
 
@@ -239,8 +243,10 @@ class SignOperation(CryptoOperation):
     output_help = ('Name of the file to write the signature to. Written '
                    'to stdout if this option is not present.')
 
-    def perform_operation(self, indata, priv_key, cli_args):
+    def perform_operation(self, indata: bytes, priv_key: rsa.key.AbstractKey,
+            cli_args: Indexable):
         """Signs files."""
+        assert isinstance(priv_key, rsa.key.PrivateKey)
 
         hash_method = cli_args[1]
         if hash_method not in HASH_METHODS:
@@ -264,8 +270,10 @@ class VerifyOperation(CryptoOperation):
     expected_cli_args = 2
     has_output = False
 
-    def perform_operation(self, indata, pub_key, cli_args):
+    def perform_operation(self, indata: bytes, pub_key: rsa.key.AbstractKey,
+            cli_args: Indexable):
         """Verifies files."""
+        assert isinstance(pub_key, rsa.key.PublicKey)
 
         signature_file = cli_args[1]
 

rsa/common.py

@@ -16,17 +16,18 @@
 
 """Common functionality shared by several modules."""
 
+import typing
+
 
 class NotRelativePrimeError(ValueError):
-    def __init__(self, a, b, d, msg=None):
-        super(NotRelativePrimeError, self).__init__(
-            msg or "%d and %d are not relatively prime, divider=%i" % (a, b, d))
+    def __init__(self, a, b, d, msg=''):
+        super().__init__(msg or "%d and %d are not relatively prime, divider=%i" % (a, b, d))
         self.a = a
         self.b = b
         self.d = d
 
 
-def bit_size(num):
+def bit_size(num: int) -> int:
     """
     Number of bits needed to represent a integer excluding any prefix
     0 bits.
@@ -54,7 +55,7 @@ def bit_size(num):
         raise TypeError('bit_size(num) only supports integers, not %r' % type(num))
 
 
-def byte_size(number):
+def byte_size(number: int) -> int:
     """
     Returns the number of bytes required to hold a specific long number.
 
@@ -79,7 +80,7 @@ def byte_size(number):
     return ceil_div(bit_size(number), 8)
 
 
-def ceil_div(num, div):
+def ceil_div(num: int, div: int) -> int:
     """
     Returns the ceiling function of a division between `num` and `div`.
 
@@ -103,7 +104,7 @@ def ceil_div(num, div):
     return quanta
 
 
-def extended_gcd(a, b):
+def extended_gcd(a: int, b: int) -> typing.Tuple[int, int, int]:
     """Returns a tuple (r, i, j) such that r = gcd(a, b) = ia + jb
     """
     # r = gcd(a,b) i = multiplicitive inverse of a mod b
@@ -128,7 +129,7 @@ def extended_gcd(a, b):
     return a, lx, ly  # Return only positive values
 
 
-def inverse(x, n):
+def inverse(x: int, n: int) -> int:
     """Returns the inverse of x % n under multiplication, a.k.a x^-1 (mod n)
 
     >>> inverse(7, 4)
@@ -145,7 +146,7 @@ def inverse(x, n):
     return inv
 
 
-def crt(a_values, modulo_values):
+def crt(a_values: typing.Iterable[int], modulo_values: typing.Iterable[int]) -> int:
     """Chinese Remainder Theorem.
 
     Calculates x such that x = a[i] (mod m[i]) for each i.

rsa/core.py

@@ -21,14 +21,14 @@
 """
 
 
-def assert_int(var, name):
+def assert_int(var: int, name: str):
     if isinstance(var, int):
         return
 
     raise TypeError('%s should be an integer, not %s' % (name, var.__class__))
 
 
-def encrypt_int(message, ekey, n):
+def encrypt_int(message: int, ekey: int, n: int) -> int:
     """Encrypts a message using encryption key 'ekey', working modulo n"""
 
     assert_int(message, 'message')
@@ -44,7 +44,7 @@ def encrypt_int(message, ekey, n):
     return pow(message, ekey, n)
 
 
-def decrypt_int(cyphertext, dkey, n):
+def decrypt_int(cyphertext: int, dkey: int, n: int) -> int:
     """Decrypts a cypher text using the decryption key 'dkey', working modulo n"""
 
     assert_int(cyphertext, 'cyphertext')