added more type hints

Author: manfred-kaiser

ssh_proxy_server/forwarders/tunnel.py

@@ -13,8 +13,8 @@
 from typeguard import typechecked
 
 import ssh_proxy_server
-from ssh_proxy_server.interfaces.server import ServerInterface
 if TYPE_CHECKING:
+    from ssh_proxy_server.interfaces.server import ServerInterface
     from ssh_proxy_server.session import Session
 
 
@@ -164,7 +164,7 @@ class ServerTunnelForwarder(ServerTunnelBaseForwarder):
     def __init__(
         self,
         session: 'ssh_proxy_server.session.Session',
-        server_interface: ServerInterface,
+        server_interface: 'ssh_proxy_server.interfaces.server.ServerInterface',
         destination: Optional[Tuple[str, int]]
     ) -> None:
         super(ServerTunnelBaseForwarder, self).__init__()

ssh_proxy_server/multisocket.py

@@ -1,4 +1,3 @@
-# type: ignore
 """
 Utility functions to create server sockets able to listen on both
 IPv4 and IPv6.
@@ -39,13 +38,18 @@
     Tuple,
     Optional,
     Text,
-    List
+    List,
+    Union,
+    overload
 )
 
+from typeguard import typechecked
+
 __author__ = "Giampaolo Rodola' <g.rodola [AT] gmail [DOT] com>"
 __license__ = "MIT"
 
 
+@typechecked
 def has_dual_stack(sock: Optional[socket.socket] = None) -> bool:
     """Return True if kernel allows creating a socket which is able to
     listen for both IPv4 and IPv6 connections.
@@ -64,6 +68,7 @@ def has_dual_stack(sock: Optional[socket.socket] = None) -> bool:
         return False
 
 
+@typechecked
 def create_server_sock(
     address: Tuple[Text, int],
     family: Optional[socket.AddressFamily] = None,
@@ -170,6 +175,7 @@ class MultipleSocketsListener:
     socket in the list.
     """
 
+    @typechecked
     def __init__(
         self,
         addresses: List[Tuple[Text, int]],
@@ -179,8 +185,8 @@ def __init__(
         queue_size: int = 5
     ) -> None:
         self._pollster: Optional[select.poll]
-        self._socks = []
-        self._sockmap = {}
+        self._socks: List[socket.socket] = []
+        self._sockmap: Dict[int, socket.socket] = {}
         if hasattr(select, 'poll'):
             self._pollster = select.poll()
         else:
@@ -206,13 +212,16 @@ def __init__(
             if not completed:
                 self.close()
 
-    def __enter__(self):
+    @typechecked
+    def __enter__(self) -> 'MultipleSocketsListener':
         return self
 
-    def __exit__(self, *args):
+    @typechecked
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         self.close()
 
-    def __repr__(self):
+    @typechecked
+    def __repr__(self) -> Text:
         addrs = []
         for sock in self._socks:
             try:
@@ -221,72 +230,101 @@ def __repr__(self):
                 addrs.append(())
         return '<%s (%r) at %#x>' % (self.__class__.__name__, addrs, id(self))
 
-    def _poll(self):
+    @typechecked
+    def _poll(self) -> Optional[Any]:
         """Return the first readable fd."""
+        fds_select: Optional[Tuple[List[Any], List[Any], List[Any]]] = None
+        fds_poll: Optional[List[Tuple[int, int]]] = None
         timeout = self.gettimeout()
         if self._pollster is None:
-            fds = select.select(self._sockmap.keys(), [], [], timeout)
-            if timeout and fds == ([], [], []):
-                raise socket.timeout('timed out')
+            fds_select = select.select(self._sockmap.keys(), [], [], timeout)
+            if timeout and fds_select == ([], [], []):
+                raise TimeoutError('timed out')
         else:
             if timeout is not None:
                 timeout *= 1000
-            fds = self._pollster.poll(timeout)
-            if timeout and fds == []:
-                raise socket.timeout('timed out')
+            fds_poll = self._pollster.poll(timeout)
+            if timeout and fds_poll == []:
+                raise TimeoutError('timed out')
         try:
-            return fds[0][0]
+            if fds_select is not None:
+                return fds_select[0][0]
+            if fds_poll is not None:
+                return fds_poll[0][0]
         except IndexError:
             pass  # non-blocking socket
+        return None
 
-    def _multicall(self, name: Text, *args: Tuple[Any], **kwargs: Dict[Text, Any]) -> None:
+    @typechecked
+    def _multicall(self, name: Text, *args: Any, **kwargs: Any) -> None:
         for sock in self._socks:
             meth = getattr(sock, name)
             meth(*args, **kwargs)
 
-    def accept(self) -> None:
+    @typechecked
+    def accept(self) -> Tuple[socket.socket, Any]:
         """Accept a connection from the first socket which is ready
         to do so.
         """
         fd = self._poll()
         sock = self._sockmap[fd] if fd else self._socks[0]
         return sock.accept()
 
-    def filenos(self):
+    @typechecked
+    def filenos(self) -> List[int]:
         """Return sockets' file descriptors as a list of integers.
         This is useful with select().
         """
         return list(self._sockmap.keys())
 
 
-    def getsockname(self):
+    @typechecked
+    def getsockname(self) -> Any:
         """Return first registered socket's own address."""
         return self._socks[0].getsockname()
 
-    def getsockopt(self, level, optname, buflen=0):
+    @overload
+    def getsockopt(self, level: int, optname: int) -> int: ...
+
+    @overload
+    def getsockopt(self, level: int, optname: int, buflen: int) -> bytes: ...
+
+    @typechecked
+    def getsockopt(self, level: int, optname: int, buflen: int = 0) -> Union[int, bytes]:
         """Return first registered socket's options."""
         return self._socks[0].getsockopt(level, optname, buflen)
 
-    def gettimeout(self) -> float:
+    @typechecked
+    def gettimeout(self) -> Optional[float]:
         """Return first registered socket's timeout."""
         return self._socks[0].gettimeout()
 
+    @typechecked
     def settimeout(self, timeout: float) -> None:
         """Set timeout for all registered sockets."""
         self._multicall('settimeout', timeout)
 
-    def setblocking(self, flag):
+    @typechecked
+    def setblocking(self, flag: bool) -> None:
         """Set non/blocking mode for all registered sockets."""
         self._multicall('setblocking', flag)
 
-    def setsockopt(self, level, optname, value):
+    @overload
+    def setsockopt(self, level: int, optname: int, value: Union[int, bytes], optlen: None) -> None: ...
+    @overload
+    def setsockopt(self, level: int, optname: int, value: None, optlen: int) -> None: ...
+
+    @typechecked
+    def setsockopt(self, level: int, optname: int, value: Optional[Union[int, bytes]], optlen: Optional[int]) -> None:
         """Set option for all registered sockets."""
-        self._multicall('setsockopt', level, optname, value)
+        self._multicall('setsockopt', level, optname, value, optlen)
 
-    def shutdown(self, how) -> None:
+    @typechecked
+    def shutdown(self, how: int) -> None:
         """Shut down all registered sockets."""
         self._multicall('shutdown', how)
 
+    @typechecked
     def close(self) -> None:
         """Close all registered sockets."""
         self._multicall('close')

ssh_proxy_server/plugins/session/tcpserver.py

@@ -2,25 +2,46 @@
 import socket
 import threading
 import time
+from typing import (
+    Callable,
+    List,
+    Text,
+    Any,
+    Union,
+    Tuple,
+    Optional
+)
+import paramiko
+
+from typeguard import typechecked
 
 
 class TCPServerThread(threading.Thread):
 
-    def __init__(self, request_handler, network='127.0.0.1', port=0, run_status=True, daemon=False):
+    @typechecked
+    def __init__(
+        self,
+        request_handler: Optional[Callable[[paramiko.Channel, Tuple[Text, int]], None]] = None,
+        network: Text = '127.0.0.1',
+        port: int = 0,
+        run_status: bool = True,
+        daemon: bool = False
+    ) -> None:
         super(TCPServerThread, self).__init__()
         self.running = run_status
         self.network = network
         self.port = port
-        self.handle_request = request_handler
+        self.handle_request_callback = request_handler
         self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         if daemon:
             self.daemon = True
         self.socket.bind((self.network, self.port))
         self.network, self.port = self.socket.getsockname()
         self.socket.listen(5)
-        self.threads = []
+        self.threads: List[threading.Thread] = []
 
+    @typechecked
     def run(self) -> None:
         while self.running:
             readable = select.select([self.socket], [], [], 0.5)[0]
@@ -30,10 +51,13 @@ def run(self) -> None:
                 t.start()
             time.sleep(0.1)
 
-    def handle_request(self, client, addr):
-        pass
+    @typechecked
+    def handle_request(self, client: paramiko.Channel, addr: Tuple[Text, int]) -> None:
+        if self.handle_request_callback is not None:
+            self.handle_request_callback(client, addr)
 
-    def close(self):
+    @typechecked
+    def close(self) -> None:
         for t in self.threads:
             t.join()
         self.socket.close()