fixed some errors

Author: manfred-kaiser

ssh_proxy_server/forwarders/tunnel.py

@@ -2,10 +2,12 @@
 import select
 import threading
 import time
+from socket import socket
 from typing import (
     TYPE_CHECKING,
     Optional,
-    Tuple
+    Tuple,
+    Union
 )
 
 import paramiko
@@ -21,7 +23,7 @@
 class TunnelForwarder(threading.Thread):
 
     @typechecked
-    def __init__(self, local_ch: Optional[paramiko.Channel], remote_ch: Optional[paramiko.Channel]) -> None:
+    def __init__(self, local_ch: Optional[Union[socket, paramiko.Channel]], remote_ch: Optional[Union[socket, paramiko.Channel]]) -> None:
         super(TunnelForwarder, self).__init__()
         self.local_ch = local_ch
         self.remote_ch = remote_ch
@@ -87,7 +89,7 @@ def close(self) -> None:
             self.close_channel(self.remote_ch)
 
     @typechecked
-    def close_channel(self, channel: paramiko.Channel) -> None:
+    def close_channel(self, channel: Union[socket, paramiko.Channel]) -> None:
         if not isinstance(channel, paramiko.Channel): # socket.socket
             channel.close()
             return

ssh_proxy_server/plugins/session/tcpserver.py

@@ -21,7 +21,7 @@ class TCPServerThread(threading.Thread):
     @typechecked
     def __init__(
         self,
-        request_handler: Optional[Callable[[paramiko.Channel, Tuple[Text, int]], None]] = None,
+        request_handler: Optional[Callable[[Union[socket.socket, paramiko.Channel], Tuple[Text, int]], None]] = None,
         network: Text = '127.0.0.1',
         port: int = 0,
         run_status: bool = True,
@@ -52,7 +52,7 @@ def run(self) -> None:
             time.sleep(0.1)
 
     @typechecked
-    def handle_request(self, client: paramiko.Channel, addr: Tuple[Text, int]) -> None:
+    def handle_request(self, client: Union[socket.socket, paramiko.Channel], addr: Tuple[Text, int]) -> None:
         if self.handle_request_callback is not None:
             self.handle_request_callback(client, addr)
 

ssh_proxy_server/plugins/tunnel/injectclienttunnel.py

@@ -35,7 +35,7 @@ def __init__(
         self.destination = destination
 
     @typechecked
-    def handle_request(self, client: paramiko.Channel, addr: Optional[Tuple[str, int]]) -> None:
+    def handle_request(self, client: Union[socket.socket, paramiko.Channel], addr: Optional[Tuple[str, int]]) -> None:
         if self.session.ssh_client is None or self.session.ssh_client.transport is None:
             return
         try:

ssh_proxy_server/plugins/tunnel/injectservertunnel.py

@@ -1,5 +1,6 @@
 import logging
-from typing import TYPE_CHECKING, Optional, Tuple, Text
+from socket import socket
+from typing import TYPE_CHECKING, Optional, Tuple, Text, Union
 
 import paramiko
 from typeguard import typechecked
@@ -49,7 +50,7 @@ def __init__(
         self.tcpserver.start()
 
     @typechecked
-    def serve(self, client: paramiko.Channel, addr: Tuple[Text, int]) -> None:
+    def serve(self, client: Union[socket, paramiko.Channel], addr: Tuple[Text, int]) -> None:
         try:
             f = TunnelForwarder(
                 self.session.transport.open_channel("forwarded-tcpip", self.destination, addr),