added type hints

Author: manfred-kaiser

ssh_proxy_server/authentication.py

@@ -33,16 +33,16 @@
 def probe_host(hostname_or_ip: Text, port: int, username: Text, public_key: paramiko.pkey.PublicBlob) -> bool:
 
     @typechecked
-    def valid(self, msg: paramiko.message.Message) -> None:
+    def valid(self, msg: paramiko.message.Message) -> None:  # type: ignore
         self.auth_event.set()
         self.authenticated = True
 
     @typechecked
-    def parse_service_accept(self, m: paramiko.message.Message) -> None:
+    def parse_service_accept(self, m: paramiko.message.Message) -> None:  # type: ignore
         # https://tools.ietf.org/html/rfc4252#section-7
         service = m.get_text()
         if not (service == "ssh-userauth" and self.auth_method == "publickey"):
-            return self._parse_service_accept(m)
+            return self._parse_service_accept(m)  # type: ignore
         m = paramiko.message.Message()
         m.add_byte(paramiko.common.cMSG_USERAUTH_REQUEST)
         m.add_string(self.username)
@@ -83,16 +83,16 @@ class RemoteCredentials():
     @typechecked
     def __init__(
         self, *,
-        username: str,
-        password: Optional[str] = None,
-        key=None,
-        host: Optional[str] = None,
+        username: Text,
+        password: Optional[Text] = None,
+        key: Optional[PKey]=None,
+        host: Optional[Text] = None,
         port: Optional[int] = None
     ) -> None:
-        self.username: str = username
-        self.password: Optional[str] = password
-        self.key = key
-        self.host: Optional[str] = host
+        self.username: Text = username
+        self.password: Optional[Text] = password
+        self.key: Optional[PKey] = key
+        self.host: Optional[Text] = host
         self.port: Optional[int] = port
 
 
@@ -177,8 +177,8 @@ def __init__(self, session: 'ssh_proxy_server.session.Session') -> None:
     @typechecked
     def get_remote_host_credentials(
         self,
-        username: str,
-        password: Optional[str] = None,
+        username: Text,
+        password: Optional[Text] = None,
         key: Optional[PKey] = None
     ) -> RemoteCredentials:
         if self.session.proxyserver.transparent:
@@ -199,7 +199,7 @@ def get_remote_host_credentials(
 
     @classmethod
     @typechecked
-    def get_auth_methods(cls, host: str, port: int) -> Optional[List[str]]:
+    def get_auth_methods(cls, host: Text, port: int) -> Optional[List[Text]]:
         auth_methods = None
         t = paramiko.Transport((host, port))
         try:
@@ -218,8 +218,8 @@ def get_auth_methods(cls, host: str, port: int) -> Optional[List[str]]:
     @typechecked
     def authenticate(
         self,
-        username: Optional[str] = None,
-        password: Optional[str] = None,
+        username: Optional[Text] = None,
+        password: Optional[Text] = None,
         key: Optional[PKey] = None,
         store_credentials: bool = True
     ) -> int:

ssh_proxy_server/clients/sftp.py

@@ -1,23 +1,47 @@
 import logging
 import os
 from typing import (
-    Optional
+    TYPE_CHECKING,
+    Callable,
+    List,
+    Any,
+    Optional,
+    Tuple,
+    Union,
+    Text
 )
 
 import paramiko
+from paramiko.pkey import PKey
+from paramiko.sftp_attr import SFTPAttributes
+from paramiko.sftp_file import SFTPFile
+from typeguard import typechecked
 
-from ssh_proxy_server.clients.ssh import SSHClient
-
+import ssh_proxy_server
+from ssh_proxy_server.clients.ssh import AuthenticationMethod, SSHClient
+if TYPE_CHECKING:
+    from ssh_proxy_server.session import Session
 
 class SFTPClient(SSHClient):
 
-    def __init__(self, host, port, method, password, user, key, session):
+    @typechecked
+    def __init__(
+        self,
+        host: Text,
+        port: int,
+        method: AuthenticationMethod,
+        password: Optional[Text],
+        user: Text,
+        key: Optional[PKey],
+        session: 'ssh_proxy_server.session.Session'
+    ) -> None:
         super().__init__(host, port, method, password, user, key, session)
         self._sftp: Optional[paramiko.SFTPClient] = None
         self.subsystem_count = 0
 
     @classmethod
-    def from_client(cls, ssh_client: Optional[SSHClient]):
+    @typechecked
+    def from_client(cls, ssh_client: Optional[SSHClient]) -> Optional['SFTPClient']:
         if ssh_client is None:
             logging.error('error creating sftp client - no ssh client!')
             return None
@@ -46,35 +70,47 @@ def from_client(cls, ssh_client: Optional[SSHClient]):
             return None
 
     @property
-    def running(self):
+    def running(self) -> bool:
         return self.subsystem_count > 0
 
-    def connect(self):
+    @typechecked
+    def connect(self) -> bool:
         ret = super().connect()
         if not ret:
             return False
         if self._sftp is None:
-            return paramiko.sftp.SFTP_FAILURE
+            return False
         try:
+            if self.transport is None:
+                return False
             self._sftp = paramiko.SFTPClient.from_transport(self.transport)
             return True
         except Exception:
             logging.exception('error creating sftp client')
         return False
 
-    def chmod(self, path, mode):
+    @typechecked
+    def open(self, filename: Union[Text, bytes], mode: Text = 'r', bufsize: int = -1) -> SFTPFile:
+        if self._sftp is None:
+            raise paramiko.SFTPError("Expected handle")
+        return self._sftp.open(filename, mode, bufsize)
+
+    @typechecked
+    def chmod(self, path: Union[Text, bytes], mode: int) -> int:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         self._sftp.chmod(path, mode)
         return paramiko.sftp.SFTP_OK
 
-    def chown(self, path, uid, gid):
+    @typechecked
+    def chown(self, path: Union[Text, bytes], uid: int, gid: int) -> int:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         self._sftp.chown(path, uid, gid)
         return paramiko.sftp.SFTP_OK
 
-    def get(self, remotePath, localPath, callback=None):
+    @typechecked
+    def get(self, remotePath: Union[Text, bytes], localPath: Union[Text, bytes], callback: Optional[Callable[[int, int], Any]] = None) -> int:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         try:
@@ -83,70 +119,84 @@ def get(self, remotePath, localPath, callback=None):
         except (IOError, OSError) as ex:
             logging.error(ex)
             os.remove(localPath)
-            return paramiko.sftp.SFTP_FAILURE
+        return paramiko.sftp.SFTP_FAILURE
 
-    def listdir_attr(self, path='.'):
+    @typechecked
+    def listdir_attr(self, path: Text = '.') -> Union[int, List[SFTPAttributes]]:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         return self._sftp.listdir_attr(path)
 
-    def lstat(self, path):
+    @typechecked
+    def lstat(self, path: Union[Text, bytes]) -> Union[int, SFTPAttributes]:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         return self._sftp.lstat(path)
 
-    def mkdir(self, path, mode=511):
+    @typechecked
+    def mkdir(self, path: Union[Text, bytes], mode: int = 511) -> int:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         self._sftp.mkdir(path, mode)
         return paramiko.sftp.SFTP_OK
 
-    def put(self, localPath, remotePath, callback=None, confirm=True):
+    @typechecked
+    def put(self, localPath: Union[Text, bytes], remotePath: Union[Text, bytes], callback: Any = None, confirm: bool = True) -> None:
         raise NotImplementedError('put not implemented')
 
-    def readlink(self, path):
+    @typechecked
+    def readlink(self, path: Union[Text, bytes]) -> Union[int, Text]:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
-        return self._sftp.readlink(path)
+        return self._sftp.readlink(path) or paramiko.sftp.SFTP_FAILURE
 
-    def remove(self, path):
+    @typechecked
+    def remove(self, path: Union[Text, bytes]) -> int:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         self._sftp.remove(path)
         return paramiko.sftp.SFTP_OK
 
-    def rename(self, oldpath, newpath):
+    @typechecked
+    def rename(self, oldpath: Union[Text, bytes], newpath: Union[Text, bytes]) -> int:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         self._sftp.rename(oldpath, newpath)
         return paramiko.sftp.SFTP_OK
 
-    def rmdir(self, path):
+    @typechecked
+    def rmdir(self, path: Union[Text, bytes]) -> int:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         self._sftp.rmdir(path)
         return paramiko.sftp.SFTP_OK
 
-    def stat(self, path):
+    @typechecked
+    def stat(self, path: Union[Text, bytes]) -> Union[int, SFTPAttributes]:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         return self._sftp.stat(path)
 
-    def utime(self, path, times):
+    @typechecked
+    def utime(self, path: Union[Text, bytes], times: Tuple[float, float]) -> int:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
-        return self._sftp.utime(path, times)
+        self._sftp.utime(path, times)
+        return paramiko.sftp.SFTP_OK
 
-    def symlink(self, source, dest):
+    @typechecked
+    def symlink(self, source: Union[Text, bytes], dest: Union[Text, bytes]) -> int:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         self._sftp.symlink(source, dest)
         return paramiko.sftp.SFTP_OK
 
-    def close(self):
+    @typechecked
+    def close(self) -> int:
         if self._sftp is None:
             return paramiko.sftp.SFTP_FAILURE
         if not self.running:
             self._sftp.close()
-            self.session.sftp_channel.close()
+            if self.session.sftp_channel is not None:
+                self.session.sftp_channel.close()
         return paramiko.sftp.SFTP_OK

ssh_proxy_server/clients/ssh.py

@@ -15,6 +15,7 @@
 
 from enhancements.modules import BaseModule
 import ssh_proxy_server
+from ssh_proxy_server.forwarders.agent import AgentProxy
 from ssh_proxy_server.exceptions import NoAgentKeys, InvalidHostKey
 
 if TYPE_CHECKING:
@@ -35,7 +36,7 @@ class SSHClient(BaseSSHClient):
 
     CIPHERS = None
 
-    #@typechecked
+    @typechecked
     def __init__(
         self,
         host: Text,
@@ -46,18 +47,18 @@ def __init__(
         key: Optional[PKey],
         session: 'ssh_proxy_server.session.Session'
     ) -> None:
-        self.session = session
-        self.host = host
-        self.port = port
-        self.method = method
-        self.user = user
-        self.password = password
-        self.agent = session.agent
-        self.key = key
-        self.transport = None
-        self.connected = False
-
-    def connect(self):
+        self.session: 'ssh_proxy_server.session.Session' = session
+        self.host: Text = host
+        self.port: int = port
+        self.method: AuthenticationMethod = method
+        self.user: Text = user
+        self.password: Optional[Text] = password
+        self.agent: Optional[AgentProxy] = self.session.agent
+        self.key: Optional[PKey] = key
+        self.transport: Optional[paramiko.Transport] = None
+        self.connected: bool = False
+
+    def connect(self) -> bool:
         message = None
 
         self.transport = paramiko.Transport((self.host, self.port))
@@ -72,19 +73,20 @@ def connect(self):
             elif self.method is AuthenticationMethod.publickey:
                 self.transport.connect(username=self.user, password=self.password, pkey=self.key)
             elif self.method is AuthenticationMethod.agent:
-                keys = self.agent.get_keys()
-                if not keys:
-                    raise NoAgentKeys()
-                for k in keys:
-                    try:
-                        self.transport.connect(username=self.user, password=self.password, pkey=k)
-                        ssh_pub_key = SSHKey(f"{k.get_name()} {k.get_base64()}")
-                        ssh_pub_key.parse()
-                        logging.debug("ssh-mitm connected to remote host with username=%s, key=%s %s %sbits", self.user, k.get_name(), ssh_pub_key.hash_sha256(), ssh_pub_key.bits)
-                        break
-                    except paramiko.AuthenticationException:
-                        self.transport.close()
-                        self.transport = paramiko.Transport((self.host, self.port))
+                if self.agent is not None:
+                    keys = self.agent.get_keys()
+                    if not keys:
+                        raise NoAgentKeys()
+                    for k in keys:
+                        try:
+                            self.transport.connect(username=self.user, password=self.password, pkey=k)
+                            ssh_pub_key = SSHKey(f"{k.get_name()} {k.get_base64()}")
+                            ssh_pub_key.parse()
+                            logging.debug("ssh-mitm connected to remote host with username=%s, key=%s %s %sbits", self.user, k.get_name(), ssh_pub_key.hash_sha256(), ssh_pub_key.bits)
+                            break
+                        except paramiko.AuthenticationException:
+                            self.transport.close()
+                            self.transport = paramiko.Transport((self.host, self.port))
 
             else:
                 logging.error('authentication method "%s" not supported!', self.method.value)
@@ -108,6 +110,6 @@ def connect(self):
 
         return False
 
-    def check_host_key(self, hostname, keytype, key):
+    def check_host_key(self, hostname: Text, keytype: Text, key: PKey) -> bool:
         """checks the host key, default always returns true"""
         return True

ssh_proxy_server/forwarders/base.py

@@ -1,5 +1,6 @@
 from typing import (
-    TYPE_CHECKING
+    TYPE_CHECKING,
+    Optional
 )
 
 from enhancements.modules import BaseModule
@@ -28,7 +29,7 @@ def __init__(self, session: 'ssh_proxy_server.session.Session') -> 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: paramiko.Channel = None
+        self.channel: Optional[paramiko.Channel] = None
         self.session: 'Session' = session
 
         # pass environment variables from client to server