Updated client.py to include remote_port

parameter, added documentation to protocol.py, and
improved server.py by using PacketType for type
and adding whitespace for readability.

src/client/client.py

@@ -96,16 +96,16 @@ async def websocket_listener(websocket, target_host='localhost', target_port=22)
             else:
                 logger.error(f"Invalid connection ID: {connection_id}")
 
-async def async_main(hostport, target_port, target_host):
+async def async_main(hostport, target_port, target_host, remote_port):
     async with websockets.connect(f"ws://{hostport}") as websocket:
         await websocket.send(json.dumps({
-            "type": 'tcp_listen',
-            "payload": {
-                "port": target_port,
-                "remote_host": "localhost",
-                "remote_port": 9999
-            }
-        }))
+                "type": 'tcp_listen',
+                "payload": {
+                    "port": target_port,
+                    "remote_host": "localhost",
+                    "remote_port": remote_port
+                }
+            }))
 
         await websocket_listener(websocket, target_host=target_host, target_port=target_port)
 
@@ -115,16 +115,18 @@ def main():
     parser.add_argument("--target_host", type=str, default='localhost', help="The target host")
     parser.add_argument("--schema", type=str, default='tcp', help="The schema to use")
     parser.add_argument("--hostport", type=str, default='localhost:8765', help="The host:port to bind to")
+    parser.add_argument("--remote_port", type=int, default=22, help="The remote port to connect to")
 
     args = parser.parse_args()
 
     target_port = args.port
     schema = args.schema
     hostport = args.hostport
     target_host = args.target_host
+    remote_port = args.remote_port
 
     if schema == 'tcp':
-        asyncio.run(async_main(hostport, target_port, target_host))
+        asyncio.run(async_main(hostport, target_port, target_host, remote_port))
 
 if __name__ == "__main__":
     main()

src/common/protocol.py

@@ -13,8 +13,29 @@ class PacketType(Enum):
 
 
 class Packet:
+    """
+    This class is used to store a Packet object.
+    
+    Attributes:
+        type: The type of the packet.
+        data: The raw data to be sent over the wire. hex-encoded of binary.
+        payload: The payload of the packet.
+        size: The size of the packet.
+        connection_id: The connection ID of the packet.
+    """
 
     class Payload:
+        """
+        This class is used to store the payload of a Packet object.
+
+        Attributes:
+            data: The raw data to be sent over the wire. hex-encoded of binary.
+            port: Local client port to listen on. not used.
+            remote_host: The remote host to connect to.
+            remote_port: The remote port to connect to.
+            websocket_id: The ID of the WebSocket connection.
+            data_tunnel_mode: The data tunnel mode to use.
+        """
         def __init__(self, data: dict):
             self.data = data.get('data')
             self.port = data.get('port') # 本地监听端口

src/server/server.py

@@ -30,6 +30,7 @@ async def handler(websocket: websockets.WebSocketServerProtocol, path: str):
         async for message in websocket:
             data = json.loads(message)
             data = Packet(data)
+
             if data.type == PacketType.TCP_LISTEN:
                 # 开启TCP服务器监听指定端口
                 tcp_server = await asyncio.start_server(
@@ -39,7 +40,7 @@ async def handler(websocket: websockets.WebSocketServerProtocol, path: str):
                 )
                 # 通知客户端新的TCP服务器已建立
                 response = Packet({
-                    "type": "new_tcp_server",
+                    "type": PacketType.NEW_TCP_SERVER,
                     "payload": {
                         "remote_host": tcp_server.sockets[0].getsockname()[0],
                         "remote_port": tcp_server.sockets[0].getsockname()[1],
@@ -50,8 +51,10 @@ async def handler(websocket: websockets.WebSocketServerProtocol, path: str):
                 await websocket.send(json.dumps(response))
                 # 将TCP服务器加入到活跃的TCP连接中
                 CONNECTIONS[websocket_id]['tcp_server'].append(tcp_server)
+
             elif data.type == PacketType.TCP_DATA:
                 await tcp_server_response_handler(data, websocket)
+
     except websockets.exceptions.ConnectionClosed:
         logger.info(f'Connection closed from {websocket.remote_address}')
     except Exception as e: