diff --git a/packages/jumpstarter-driver-network/jumpstarter_driver_network/adapters/portforward.py b/packages/jumpstarter-driver-network/jumpstarter_driver_network/adapters/portforward.py index 557bb14eb..7a13e2f6a 100644 --- a/packages/jumpstarter-driver-network/jumpstarter_driver_network/adapters/portforward.py +++ b/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 diff --git a/packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py b/packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py index c3db4fa7c..808afa8c3 100644 --- a/packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py +++ b/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): diff --git a/packages/jumpstarter-driver-network/pyproject.toml b/packages/jumpstarter-driver-network/pyproject.toml index bccda4c80..e5fcb3600 100644 --- a/packages/jumpstarter-driver-network/pyproject.toml +++ b/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"] diff --git a/packages/jumpstarter/jumpstarter/common/tempfile.py b/packages/jumpstarter/jumpstarter/common/tempfile.py index 60ea8b1d0..b8aebdba3 100644 --- a/packages/jumpstarter/jumpstarter/common/tempfile.py +++ b/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) diff --git a/uv.lock b/uv.lock index a7da48243..027f0ef81 100644 --- a/uv.lock +++ b/uv.lock @@ -1359,6 +1359,7 @@ dev = [ name = "jumpstarter-driver-network" source = { editable = "packages/jumpstarter-driver-network" } dependencies = [ + { name = "asyncclick" }, { name = "fabric" }, { name = "jumpstarter" }, { name = "pexpect" }, @@ -1376,6 +1377,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "asyncclick", specifier = ">=8.1.8" }, { name = "fabric", specifier = ">=3.2.2" }, { name = "jumpstarter", editable = "packages/jumpstarter" }, { name = "pexpect", specifier = ">=4.9.0" },