added more type hints

Author: manfred-kaiser

ssh_proxy_server/authentication.py

@@ -5,6 +5,7 @@
 import socket
 
 from typing import (
+    TYPE_CHECKING,
     Optional,
     List,
     Tuple,
@@ -20,9 +21,12 @@
 from sshpubkeys import SSHKey  # type: ignore
 from typeguard import typechecked
 
+import ssh_proxy_server
 from ssh_proxy_server.clients.ssh import SSHClient, AuthenticationMethod
 from ssh_proxy_server.exceptions import MissingHostException
-from ssh_proxy_server.session import Session
+
+if TYPE_CHECKING:
+    from ssh_proxy_server.session import Session
 
 
 @typechecked
@@ -62,8 +66,8 @@ def parse_service_accept(self, m: paramiko.message.Message) -> None:
 
         # For compatibility with paramiko, we need to generate a random private key and replace
         # the public key with our data.
-        key = paramiko.RSAKey.generate(2048)
-        key.public_blob = public_key  # type: ignore
+        key: PKey = paramiko.RSAKey.generate(2048)
+        key.public_blob = public_key
         transport.auth_publickey(username, key)
         valid_key = True
     except paramiko.ssh_exception.AuthenticationException:
@@ -166,7 +170,7 @@ def parser_arguments(cls) -> None:
         )
 
     @typechecked
-    def __init__(self, session: Session) -> None:
+    def __init__(self, session: 'ssh_proxy_server.session.Session') -> None:
         super().__init__()
         self.session = session
 
@@ -223,7 +227,7 @@ def authenticate(
             self.session.username_provided = username
             self.session.password_provided = password
         if username:
-            remote_credentials = self.get_remote_host_credentials(username, password, key)
+            remote_credentials: RemoteCredentials = self.get_remote_host_credentials(username, password, key)
             self.session.username = remote_credentials.username
             self.session.password = remote_credentials.password
             self.session.key = remote_credentials.key

ssh_proxy_server/cli.py

@@ -19,7 +19,7 @@
     Authenticator,
     AuthenticatorPassThrough
 )
-from ssh_proxy_server.interfaces import (
+from ssh_proxy_server.interfaces.server import (
     BaseServerInterface,
     ServerInterface
 )

ssh_proxy_server/clients/sftp.py

@@ -1,5 +1,8 @@
 import logging
 import os
+from typing import (
+    Optional
+)
 
 import paramiko
 
@@ -14,7 +17,7 @@ def __init__(self, host, port, method, password, user, key, session):
         self.subsystem_count = 0
 
     @classmethod
-    def from_client(cls, ssh_client):
+    def from_client(cls, ssh_client: Optional[SSHClient]):
         if ssh_client is None:
             logging.error('error creating sftp client - no ssh client!')
             return None
@@ -32,6 +35,9 @@ def from_client(cls, ssh_client):
                 ssh_client.key,
                 ssh_client.session
             )
+            if ssh_client.transport is None:
+                logging.debug("ssh_client does not have a transport")
+                return None
             sftp._sftp = paramiko.SFTPClient.from_transport(ssh_client.transport)
             sftp.connected = True
             return sftp

ssh_proxy_server/clients/ssh.py

@@ -3,8 +3,10 @@
 
 from typing import (
     TYPE_CHECKING,
+    Optional,
     Text
 )
+from paramiko.pkey import PKey
 
 from typeguard import typechecked
 import paramiko
@@ -34,7 +36,16 @@ class SSHClient(BaseSSHClient):
     CIPHERS = None
 
     #@typechecked
-    def __init__(self, host, port, method, password, user, key, session: 'ssh_proxy_server.session.Session') -> None:
+    def __init__(
+        self,
+        host: Text,
+        port: int,
+        method: AuthenticationMethod,
+        password: Optional[Text],
+        user: Text,
+        key: Optional[PKey],
+        session: 'ssh_proxy_server.session.Session'
+    ) -> None:
         self.session = session
         self.host = host
         self.port = port

ssh_proxy_server/exceptions.py

@@ -16,3 +16,7 @@ class MissingHostException(Exception):
 
 class KeyGenerationError(Exception):
     pass
+
+
+class MissingClient(Exception):
+    pass

ssh_proxy_server/forwarders/agent.py

@@ -1,12 +1,18 @@
 import time
+from typing import Tuple, List, Union
 
-from paramiko.agent import Agent, AgentServerProxy, AgentClientProxy
+from paramiko.agent import Agent, AgentKey, AgentServerProxy, AgentClientProxy
+from paramiko.transport import Transport
+from paramiko.channel import Channel
 import os
 
+from typeguard import typechecked
+
 class AgentProxy(object):
 
-    def __init__(self, transport) -> None:
-        self.agents = []
+    @typechecked
+    def __init__(self, transport: Transport) -> None:
+        self.agents: List[Union[Agent, AgentClientProxy]] = []
         self.transport = transport
         a = AgentServerProxy(self.transport)
         os.environ.update(a.get_env())
@@ -18,18 +24,22 @@ def __init__(self, transport) -> None:
         # agent is still sending over the channel
         # agent.close()
 
-    def get_keys(self):
+    @typechecked
+    def get_keys(self) -> Tuple[AgentKey, ...]:
         return self.keys
 
-    def forward_agent(self, chanClient):
-        chanClient.request_forward_agent(self._forward_agent_handler)
+    @typechecked
+    def forward_agent(self, chanClient: Channel) -> bool:
+        return chanClient.request_forward_agent(self._forward_agent_handler)
 
-    def _forward_agent_handler(self, chanRemote):
+    @typechecked
+    def _forward_agent_handler(self, chanRemote: Channel) -> None:
         agent = AgentServerProxy(self.transport)
         os.environ.update(agent.get_env())
         time.sleep(0.1)
         self.agents.append(AgentClientProxy(chanRemote))
 
-    def close(self):
+    @typechecked
+    def close(self) -> None:
         for a in self.agents:
             a.close()

ssh_proxy_server/forwarders/base.py

@@ -1,6 +1,15 @@
-import logging
+from typing import (
+    TYPE_CHECKING
+)
 
 from enhancements.modules import BaseModule
+import paramiko
+from typeguard import typechecked
+
+import ssh_proxy_server
+from ssh_proxy_server.exceptions import MissingClient
+if TYPE_CHECKING:
+    from ssh_proxy_server.session import Session
 
 
 class BaseForwarder(BaseModule):
@@ -11,29 +20,35 @@ class BaseForwarder(BaseModule):
     # Slow file transmission
     BUF_LEN = 65536*100
 
-    def __init__(self, session):
+    @typechecked
+    def __init__(self, session: 'ssh_proxy_server.session.Session') -> None:
         super().__init__()
-        self.server_channel = session.ssh_client.transport.open_session()
+        if session.ssh_client is None or session.ssh_client.transport is None:
+            raise MissingClient("session.ssh_client is None")
+        self.server_channel: paramiko.Channel = session.ssh_client.transport.open_session()
         if session.agent is not None:
             session.agent.forward_agent(self.server_channel)
-        self.channel = None
-        self.session = session
+        self.channel: paramiko.Channel = None
+        self.session: 'Session' = session
 
         # pass environment variables from client to server
         for env_name, env_value in self.session.env_requests.items():
             self.server_channel.set_environment_variable(env_name, env_value)
 
-    def forward(self):
+    @typechecked
+    def forward(self) -> None:
         raise NotImplementedError
 
-    def close_session(self, channel):
+    @typechecked
+    def close_session(self, channel: paramiko.Channel) -> None:
         channel.lock.acquire()
         if not channel.closed:
             channel.lock.release()
             channel.close()
         if channel.lock.locked():
             channel.lock.release()
 
-    def _closed(self, channel):
+    @typechecked
+    def _closed(self, channel: paramiko.Channel) -> bool:
         #return channel.closed or channel.eof_received or channel.eof_sent or not channel.active
         return channel.closed or not channel.active

ssh_proxy_server/forwarders/sftp.py

@@ -1,24 +1,41 @@
 import logging
+from typing import (
+    TYPE_CHECKING,
+    Optional,
+    Union, Type
+)
+
 import paramiko
 from enhancements.modules import BaseModule
+from typeguard import typechecked
+
+import ssh_proxy_server
+from ssh_proxy_server.interfaces.sftp import BaseSFTPServerInterface
+
+if TYPE_CHECKING:
+    from ssh_proxy_server.session import Session
 
 
 class SFTPHandlerBasePlugin(BaseModule):
 
-    def __init__(self, sftp, filename):
+    @typechecked
+    def __init__(self, sftp, filename) -> None:
         super().__init__()
         self.filename = filename
         self.sftp = sftp
 
     @classmethod
-    def get_interface(cls):
+    @typechecked
+    def get_interface(cls) -> Optional[Type[BaseSFTPServerInterface]]:
         return None
 
     @classmethod
-    def get_file_handle(cls):
-        return None
+    @typechecked
+    def get_file_handle(cls) -> Type['SFTPBaseHandle']:
+        return SFTPBaseHandle
 
-    def close(self):
+    @typechecked
+    def close(self) -> None:
         pass
 
     def handle_data(self, data, *, offset=None, length=None):
@@ -32,24 +49,34 @@ class SFTPHandlerPlugin(SFTPHandlerBasePlugin):
 
 class SFTPBaseHandle(paramiko.SFTPHandle):
 
-    def __init__(self, session, plugin, filename, flags=0):
+    @typechecked
+    def __init__(
+        self, session: 'ssh_proxy_server.session.Session', plugin, filename, flags: int = 0
+    ) -> None:
         super().__init__(flags)
         self.session = session
         self.plugin = plugin(self, filename)
-        self.writefile = None
-        self.readfile = None
+        self.writefile: Optional[paramiko.sftp_file.SFTPFile] = None
+        self.readfile: Optional[paramiko.sftp_file.SFTPFile] = None
 
-    def close(self):
+    @typechecked
+    def close(self) -> None:
         super().close()
         self.plugin.close()
 
-    def read(self, offset, length):
+    @typechecked
+    def read(self, offset, length) -> Union[bytes, int]:
         logging.debug("R_OFFSET: %s", offset)
+        if self.readfile is None:
+            return paramiko.sftp.SFTP_FAILURE
         data = self.readfile.read(length)
         return self.plugin.handle_data(data, length=length)
 
-    def write(self, offset, data):
+    @typechecked
+    def write(self, offset: int, data: bytes) -> int:
         logging.debug("W_OFFSET: %s", offset)
         data = self.plugin.handle_data(data, offset=offset)
+        if self.writefile is None:
+            return paramiko.sftp.SFTP_FAILURE
         self.writefile.write(data)
-        return paramiko.SFTP_OK
+        return paramiko.sftp.SFTP_OK