Repository: twisted/twisted

Issue:
twisted.conch.ssh.factory.SSHFactory.getPrimes and its overrides are not type annotated
I wrote some annotations when working on reviewing #10204 but they were kind of out of scope for that issue.


Discussion:

Resolving patch:
diff --git a/src/twisted/conch/openssh_compat/factory.py b/src/twisted/conch/openssh_compat/factory.py
--- a/src/twisted/conch/openssh_compat/factory.py
+++ b/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,7 +67,7 @@ def getPrivateKeys(self):
                     privateKeys[key.sshType()] = key
         return privateKeys
 
-    def getPrimes(self):
+    def getPrimes(self) -> Optional[Dict[int, List[Tuple[int, int]]]]:
         try:
             return primes.parseModuliFile(self.moduliRoot + "/moduli")
         except OSError:
diff --git a/src/twisted/conch/openssh_compat/primes.py b/src/twisted/conch/openssh_compat/primes.py
--- a/src/twisted/conch/openssh_compat/primes.py
+++ b/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))
diff --git a/src/twisted/conch/ssh/factory.py b/src/twisted/conch/ssh/factory.py
--- a/src/twisted/conch/ssh/factory.py
+++ b/src/twisted/conch/ssh/factory.py
@@ -12,6 +12,7 @@
 
 import random
 from itertools import chain
+from typing import Dict, List, Optional, 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: Optional[Dict[int, List[Tuple[int, int]]]]
+
     _log = Logger()
     protocol = transport.SSHServerTransport
 
@@ -32,7 +35,7 @@ class SSHFactory(protocol.Factory):
         b"ssh-connection": connection.SSHConnection,
     }
 
-    def startFactory(self):
+    def startFactory(self) -> None:
         """
         Check for public and private keys.
         """
@@ -93,24 +96,24 @@ def getPrivateKeys(self):
         """
         raise NotImplementedError("getPrivateKeys unimplemented")
 
-    def getPrimes(self):
+    def getPrimes(self) -> Optional[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
-        of tuple of (generator, prime).
-
-        @rtype: L{dict}
+        primes to use.  Returns a dictionary mapping number of bits to lists of
+        tuple of (generator, prime).
         """
 
-    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)
+
+        assert self.primes is not None, "Factory should have been started by now."
+        primesKeys = sorted(self.primes.keys(), key=keyfunc)
         realBits = primesKeys[0]
         return random.choice(self.primes[realBits])
 


Test patch:
diff --git a/src/twisted/conch/test/test_openssh_compat.py b/src/twisted/conch/test/test_openssh_compat.py
--- a/src/twisted/conch/test/test_openssh_compat.py
+++ b/src/twisted/conch/test/test_openssh_compat.py
@@ -117,7 +117,7 @@ def seteuid(euid):
         self.assertEqual(self.mockos.seteuidCalls, [0, os.geteuid()])
         self.assertEqual(self.mockos.setegidCalls, [0, os.getegid()])
 
-    def test_getPrimes(self):
+    def test_getPrimes(self) -> None:
         """
         L{OpenSSHFactory.getPrimes} should return the available primes
         in the moduli directory.
diff --git a/src/twisted/conch/test/test_ssh.py b/src/twisted/conch/test/test_ssh.py
--- a/src/twisted/conch/test/test_ssh.py
+++ b/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.
diff --git a/src/twisted/conch/test/test_transport.py b/src/twisted/conch/test/test_transport.py
--- a/src/twisted/conch/test/test_transport.py
+++ b/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):