Merge pull request #393 from jumpstarter-dev/portforward-cli

Implement forward-{tcp,port} command on network client

Author: NickCao

packages/jumpstarter-driver-network/jumpstarter_driver_network/adapters/portforward.py

@@ -1,5 +1,6 @@
 from contextlib import asynccontextmanager
 from functools import partial
+from os import PathLike
 
 from jumpstarter.client import DriverClient
 from jumpstarter.client.adapters import blocking
@@ -37,6 +38,7 @@ async def UnixPortforwardAdapter(
     *,
     client: DriverClient,
     method: str = "connect",
+    path: PathLike | None = None,
 ):
-    async with TemporaryUnixListener(partial(handler, client, method)) as addr:
+    async with TemporaryUnixListener(partial(handler, client, method), path=path) as addr:
         yield addr

packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py

@@ -1,12 +1,65 @@
 from contextlib import AbstractContextManager
+from ipaddress import IPv6Address, ip_address
+from threading import Event
 
-from .adapters import DbusAdapter
+import asyncclick as click
+
+from .adapters import DbusAdapter, TcpPortforwardAdapter, UnixPortforwardAdapter
 from .driver import DbusNetwork
 from jumpstarter.client import DriverClient
 
 
 class NetworkClient(DriverClient):
-    pass
+    def cli(self):
+        @click.group
+        def base():
+            """Generic Network Connection"""
+            pass
+
+        @base.command()
+        @click.option("--address", default="localhost", show_default=True)
+        @click.argument("port", type=int)
+        def forward_tcp(address: str, port: int):
+            """
+            Forward local TCP port to remote network
+
+            PORT is the TCP port to listen on.
+            """
+
+            with TcpPortforwardAdapter(
+                client=self,
+                local_host=address,
+                local_port=port,
+            ) as addr:
+                host = ip_address(addr[0])
+                port = addr[1]
+                match host:
+                    case IPv6Address():
+                        click.echo("[{}]:{}".format(host, port))
+                    case _:
+                        click.echo("{}:{}".format(host, port))
+
+                Event().wait()
+
+        @base.command()
+        @click.argument("path", type=click.Path(), required=False)
+        def forward_unix(path: str | None):
+            """
+            Forward local Unix domain socket to remote network
+
+            PATH is the path of the Unix domain socket to listen on,
+            defaults to a random path under $XDG_RUNTIME_DIR.
+            """
+
+            with UnixPortforwardAdapter(
+                client=self,
+                path=path,
+            ) as addr:
+                click.echo(addr)
+
+                Event().wait()
+
+        return base
 
 
 class DbusNetworkClient(NetworkClient, AbstractContextManager):

packages/jumpstarter-driver-network/pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
   "pexpect>=4.9.0",
   "fabric>=3.2.2",
   "wsproto>=1.2.0",
+  "asyncclick>=8.1.8",
 ]
 
 [project.entry-points."jumpstarter.drivers"]

packages/jumpstarter/jumpstarter/common/tempfile.py

@@ -1,4 +1,5 @@
-from contextlib import asynccontextmanager, contextmanager
+from contextlib import asynccontextmanager, contextmanager, nullcontext
+from os import PathLike
 from pathlib import Path
 from socket import AddressFamily
 from tempfile import TemporaryDirectory
@@ -15,8 +16,13 @@ def TemporarySocket():
 
 
 @asynccontextmanager
-async def TemporaryUnixListener(handler):
-    with TemporarySocket() as path:
+async def TemporaryUnixListener(handler, path: PathLike | None = None):
+    if path is not None:
+        cm = nullcontext(path)
+    else:
+        cm = TemporarySocket()
+
+    with cm as path:
         async with await create_unix_listener(path) as listener:
             async with create_task_group() as tg:
                 tg.start_soon(listener.serve, handler, tg)