annotate twisted.conch.ssh.factory.SSHFactory.getPrimes

src/twisted/conch/openssh_compat/factory.py

@@ -9,6 +9,7 @@
 
 import errno
 import os
+from typing import Dict, List, Optional, Tuple
 
 from twisted.conch.openssh_compat import primes
 from twisted.conch.ssh import common, factory, keys
@@ -66,8 +67,14 @@ def getPrivateKeys(self):
                     privateKeys[key.sshType()] = key
         return privateKeys
 
-    def getPrimes(self):
+    def getPrimes(
+        self,
+    ) -> Optional[Dict[int, List[Tuple[int, int]]]]:  # type:ignore[override]
         try:
             return primes.parseModuliFile(self.moduliRoot + "/moduli")
         except OSError:
+            # As seen in the type:ignore[override] above, this historically
+            # handled errors by (invalidly) returning None, and we should
+            # probably fix this at some point, since callers don't need to
+            # handle None given the contract of the superclass.
             return None

src/twisted/conch/openssh_compat/primes.py

@@ -10,18 +10,21 @@
 """
 
 
-def parseModuliFile(filename):
+from typing import Dict, List, Tuple
+
+
+def parseModuliFile(filename: str) -> Dict[int, List[Tuple[int, int]]]:
     with open(filename) as f:
         lines = f.readlines()
-    primes = {}
+    primes: Dict[int, List[Tuple[int, int]]] = {}
     for l in lines:
         l = l.strip()
         if not l or l[0] == "#":
             continue
-        tim, typ, tst, tri, size, gen, mod = l.split()
-        size = int(size) + 1
-        gen = int(gen)
-        mod = int(mod, 16)
+        tim, typ, tst, tri, sizestr, genstr, modstr = l.split()
+        size = int(sizestr) + 1
+        gen = int(genstr)
+        mod = int(modstr, 16)
         if size not in primes:
             primes[size] = []
         primes[size].append((gen, mod))

src/twisted/conch/ssh/factory.py

@@ -12,6 +12,7 @@
 
 import random
 from itertools import chain
+from typing import Dict, List, Tuple
 
 from twisted.conch import error
 from twisted.conch.ssh import _kex, connection, transport, userauth
@@ -24,6 +25,8 @@ class SSHFactory(protocol.Factory):
     A Factory for SSH servers.
     """
 
+    primes: Dict[int, List[Tuple[int, int]]]
+
     _log = Logger()
     protocol = transport.SSHServerTransport
 
@@ -93,7 +96,7 @@ def getPrivateKeys(self):
         """
         raise NotImplementedError("getPrivateKeys unimplemented")
 
-    def getPrimes(self):
+    def getPrimes(self) -> Dict[int, List[Tuple[int, int]]]:
         """
         Called when the factory is started to get Diffie-Hellman generators and
         primes to use.  Returns a dictionary mapping number of bits to lists
@@ -102,15 +105,16 @@ def getPrimes(self):
         @rtype: L{dict}
         """
 
-    def getDHPrime(self, bits):
+    def getDHPrime(self, bits: int) -> Tuple[int, int]:
         """
         Return a tuple of (g, p) for a Diffe-Hellman process, with p being as
-        close to bits bits as possible.
-
-        @type bits: L{int}
-        @rtype:     L{tuple}
+        close to C{bits} bits as possible.
         """
-        primesKeys = sorted(self.primes.keys(), key=lambda i: abs(i - bits))
+
+        def keyfunc(i: int) -> int:
+            return abs(i - bits)
+
+        primesKeys = sorted(self.primes.keys(), key=keyfunc)
         realBits = primesKeys[0]
         return random.choice(self.primes[realBits])
 

src/twisted/conch/test/test_ssh.py

@@ -8,6 +8,7 @@
 
 import struct
 from itertools import chain
+from typing import Dict, List, Tuple
 
 from twisted.conch.test.keydata import (
     privateDSA_openssh,
@@ -375,7 +376,7 @@ def getPrivateKeys(self):
                 b"ssh-dss": keys.Key.fromString(privateDSA_openssh),
             }
 
-        def getPrimes(self):
+        def getPrimes(self) -> Dict[int, List[Tuple[int, int]]]:
             """
             Diffie-Hellman primes that can be used for the
             diffie-hellman-group-exchange-sha1 key exchange.

src/twisted/conch/test/test_transport.py

@@ -12,7 +12,7 @@
 import struct
 import types
 from hashlib import md5, sha1, sha256, sha384, sha512
-from typing import Optional, Type
+from typing import Dict, List, Optional, Tuple, Type
 
 from twisted import __version__ as twisted_version
 from twisted.conch.error import ConchError
@@ -292,7 +292,7 @@ def getPrivateKeys(self):
             b"ssh-dsa": keys.Key.fromString(keydata.privateDSA_openssh),
         }
 
-    def getPrimes(self):
+    def getPrimes(self) -> Dict[int, List[Tuple[int, int]]]:
         """
         Diffie-Hellman primes that can be used for key exchange algorithms
         that use group exchange to establish a prime / generator group.
@@ -307,7 +307,7 @@ def getPrimes(self):
         # tests.
         # See OpenSSHFactory.getPrimes.
         group14 = _kex.getDHGeneratorAndPrime(b"diffie-hellman-group14-sha1")
-        return {2048: (group14,), 4096: ((5, 7),)}
+        return {2048: [group14], 4096: [(5, 7)]}
 
 
 class MockOldFactoryPublicKeys(MockFactory):