#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""
SSH tunnel implementations for the Airflow SSH provider.
This module provides ``SSHTunnel`` (sync, paramiko-based) and ``AsyncSSHTunnel``
(async, asyncssh-based) as replacements for the removed ``sshtunnel.SSHTunnelForwarder``.
**SSHTunnel** reuses an already-connected ``paramiko.SSHClient`` from
``SSHHook.get_conn()``, so all authentication and proxy configuration is inherited
automatically. It binds a local TCP socket and forwards connections to a remote
host/port through the SSH transport using ``open_channel('direct-tcpip', ...)``.
**AsyncSSHTunnel** wraps ``asyncssh.forward_local_port()`` and is intended for use
with ``SSHHookAsync``.
Migration from ``sshtunnel.SSHTunnelForwarder``
-----------------------------------------------
Before::
from sshtunnel import SSHTunnelForwarder
tunnel = hook.get_tunnel(remote_port=5432)
tunnel.start()
# use tunnel.local_bind_port
tunnel.stop()
After::
with hook.get_tunnel(remote_port=5432) as tunnel:
# use tunnel.local_bind_port
The ``.start()`` / ``.stop()`` methods still exist but emit deprecation warnings.
Use the context manager interface instead.
"""
from __future__ import annotations
import logging
import socket
import threading
import warnings
from select import select
from typing import TYPE_CHECKING
from airflow.exceptions import AirflowProviderDeprecationWarning
if TYPE_CHECKING:
import asyncssh
import paramiko
from airflow.sdk.types import Logger
[docs]
log = logging.getLogger(__name__)
# Attributes that existed on SSHTunnelForwarder but not on SSHTunnel.
# Used by __getattr__ to provide a helpful migration message.
_SSHTUNNELFORWARDER_ATTRS = frozenset(
{
"tunnel_is_up",
"skip_tunnel_checkup",
"ssh_host",
"ssh_port",
"ssh_username",
"ssh_password",
"ssh_pkey",
"ssh_proxy",
"local_bind_address",
"local_bind_addresses",
"local_bind_host",
"local_bind_hosts",
"remote_bind_address",
"remote_bind_addresses",
"tunnel_bindings",
"is_alive",
"raise_fwd_exc",
"daemon_forward_servers",
"daemon_transport",
}
)
[docs]
class SSHTunnel:
"""
Local port-forwarding tunnel over an existing paramiko SSH connection.
This replaces ``sshtunnel.SSHTunnelForwarder`` by using the SSH client's
transport directly via ``open_channel('direct-tcpip', ...)``.
The recommended usage is as a context manager::
client = hook.get_conn()
with SSHTunnel(client, "dbhost", 5432) as tunnel:
connect_to_db("localhost", tunnel.local_bind_port)
:param ssh_client: An already-connected ``paramiko.SSHClient``.
:param remote_host: The destination host to forward to (from the SSH server's perspective).
:param remote_port: The destination port to forward to.
:param local_port: Local port to bind. ``None`` means an OS-assigned ephemeral port.
:param logger: Optional logger instance. Falls back to the module logger.
"""
def __init__(
self,
ssh_client: paramiko.SSHClient,
remote_host: str,
remote_port: int,
local_port: int | None = None,
logger: logging.Logger | Logger | None = None,
) -> None:
self._ssh_client = ssh_client
self._remote_host = remote_host
self._remote_port = remote_port
self._logger: logging.Logger | Logger = logger or log
self._server_socket: socket.socket | None = None
self._thread: threading.Thread | None = None
# Self-pipe for waking the select loop on shutdown
self._shutdown_r: socket.socket | None = None
self._shutdown_w: socket.socket | None = None
self._running = False
# Bind the listening socket eagerly so local_bind_port is available
# before entering the context manager or calling start().
self._server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
try:
self._server_socket.bind(("localhost", local_port or 0))
self._server_socket.listen(5)
except OSError:
self._server_socket.close()
raise
self._server_socket.setblocking(False)
# -- Public properties ---------------------------------------------------
@property
[docs]
def local_bind_port(self) -> int:
"""Return the local port the tunnel is listening on."""
if self._server_socket is None:
raise RuntimeError("Tunnel socket is not bound")
return self._server_socket.getsockname()[1]
@property
[docs]
def local_bind_address(self) -> tuple[str, int]:
"""Return ``('localhost', <port>)`` — the local address the tunnel is listening on."""
return ("localhost", self.local_bind_port)
# -- Context manager -----------------------------------------------------
[docs]
def __enter__(self) -> SSHTunnel:
self._start_forwarding()
return self
[docs]
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
self._stop_forwarding()
# -- Deprecated start/stop -----------------------------------------------
[docs]
def start(self) -> None:
"""Start the tunnel. **Deprecated** — use the context manager interface instead."""
warnings.warn(
"SSHTunnel.start() is deprecated. Use the context manager interface: "
"`with hook.get_tunnel(...) as tunnel:`",
AirflowProviderDeprecationWarning,
stacklevel=2,
)
self._start_forwarding()
[docs]
def stop(self) -> None:
"""Stop the tunnel. **Deprecated** — use the context manager interface instead."""
warnings.warn(
"SSHTunnel.stop() is deprecated. Use the context manager interface: "
"`with hook.get_tunnel(...) as tunnel:`",
AirflowProviderDeprecationWarning,
stacklevel=2,
)
self._stop_forwarding()
# -- Migration helper ----------------------------------------------------
[docs]
def __getattr__(self, name: str):
if name in _SSHTUNNELFORWARDER_ATTRS:
raise AttributeError(
f"'{type(self).__name__}' has no attribute '{name}'. "
f"sshtunnel.SSHTunnelForwarder has been replaced by SSHTunnel. "
f"Use the context manager interface and .local_bind_port instead."
)
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
# -- Internal ------------------------------------------------------------
def _start_forwarding(self) -> None:
"""Start the forwarding thread."""
if self._running:
return
self._shutdown_r, self._shutdown_w = socket.socketpair()
self._running = True
self._thread = threading.Thread(target=self._serve_forever, daemon=True)
self._thread.start()
def _stop_forwarding(self) -> None:
"""Signal the forwarding thread to stop and wait for it."""
if not self._running:
return
self._running = False
# Wake the select loop
if self._shutdown_w is not None:
try:
self._shutdown_w.send(b"\x00")
except OSError:
pass
if self._thread is not None:
self._thread.join(timeout=5)
self._thread = None
# Close the shutdown pair
for sock in (self._shutdown_r, self._shutdown_w):
if sock is not None:
try:
sock.close()
except OSError:
pass
self._shutdown_r = None
self._shutdown_w = None
# Close the server socket
if self._server_socket is not None:
try:
self._server_socket.close()
except OSError:
pass
self._server_socket = None
def _serve_forever(self) -> None:
"""Accept connections on the local socket and forward them through SSH."""
if self._server_socket is None or self._shutdown_r is None:
return
server_socket = self._server_socket
shutdown_r = self._shutdown_r
active_channels: list[tuple[socket.socket, paramiko.Channel]] = []
try:
while self._running:
read_fds: list[socket.socket | paramiko.Channel] = [server_socket, shutdown_r]
for local_sock, chan in active_channels:
read_fds.append(local_sock)
read_fds.append(chan)
try:
readable, _, _ = select(read_fds, [], [], 1.0)
except (OSError, ValueError):
break
for fd in readable:
if fd is shutdown_r:
return
if fd is server_socket:
self._accept_connection(active_channels)
else:
self._forward_data(fd, active_channels)
# Clean up closed channels
active_channels = [(s, c) for s, c in active_channels if not (s.fileno() == -1 or c.closed)]
finally:
for local_sock, chan in active_channels:
self._close_pair(local_sock, chan)
def _accept_connection(self, active_channels: list[tuple[socket.socket, paramiko.Channel]]) -> None:
"""Accept a new local connection and open an SSH channel for it."""
if self._server_socket is None:
return
try:
client_sock, addr = self._server_socket.accept()
except OSError:
return
transport = self._ssh_client.get_transport()
if transport is None or not transport.is_active():
self._logger.warning("SSH transport is not active, rejecting connection from %s", addr)
client_sock.close()
return
try:
channel = transport.open_channel(
"direct-tcpip",
(self._remote_host, self._remote_port),
addr,
)
except Exception:
self._logger.warning(
"Failed to open SSH channel to %s:%s",
self._remote_host,
self._remote_port,
exc_info=True,
)
client_sock.close()
return
if channel is None:
self._logger.warning("SSH channel request was rejected")
client_sock.close()
return
active_channels.append((client_sock, channel))
def _forward_data(
self,
fd: socket.socket | paramiko.Channel,
active_channels: list[tuple[socket.socket, paramiko.Channel]],
) -> None:
"""Forward data between a local socket and its paired SSH channel."""
for local_sock, chan in active_channels:
if fd is local_sock:
try:
data = local_sock.recv(16384)
except OSError:
data = b""
if not data:
self._close_pair(local_sock, chan)
return
try:
chan.sendall(data)
except OSError:
self._logger.warning("Error sending data to SSH channel, closing connection")
self._close_pair(local_sock, chan)
return
if fd is chan:
try:
data = chan.recv(16384)
except OSError:
data = b""
if not data:
self._close_pair(local_sock, chan)
return
try:
local_sock.sendall(data)
except OSError:
self._logger.warning("Error sending data to local socket, closing connection")
self._close_pair(local_sock, chan)
return
@staticmethod
def _close_pair(local_sock: socket.socket, chan: paramiko.Channel) -> None:
"""Close both ends of a forwarded connection."""
for closeable in (chan, local_sock):
try:
closeable.close()
except OSError:
pass
[docs]
class AsyncSSHTunnel:
"""
Async local port-forwarding tunnel over an asyncssh SSH connection.
This wraps ``asyncssh.SSHClientConnection.forward_local_port()`` and is
intended for use with ``SSHHookAsync``.
Usage::
async with await hook.get_tunnel(remote_port=5432) as tunnel:
connect_to_db("localhost", tunnel.local_bind_port)
On exit, both the forwarding listener and the underlying SSH connection
are closed.
:param ssh_conn: An ``asyncssh.SSHClientConnection``.
:param remote_host: The destination host to forward to.
:param remote_port: The destination port to forward to.
:param local_port: Local port to bind. ``None`` means an OS-assigned ephemeral port.
"""
def __init__(
self,
ssh_conn: asyncssh.SSHClientConnection,
remote_host: str,
remote_port: int,
local_port: int | None = None,
) -> None:
self._ssh_conn = ssh_conn
self._remote_host = remote_host
self._remote_port = remote_port
self._local_port = local_port or 0
self._listener: asyncssh.SSHListener | None = None
@property
[docs]
def local_bind_port(self) -> int:
"""Return the local port the tunnel is listening on."""
if self._listener is None:
raise RuntimeError("Tunnel is not started. Use `async with` to start it.")
return self._listener.get_port()
[docs]
async def __aenter__(self) -> AsyncSSHTunnel:
try:
self._listener = await self._ssh_conn.forward_local_port(
"localhost",
self._local_port,
self._remote_host,
self._remote_port,
)
except BaseException:
self._ssh_conn.close()
await self._ssh_conn.wait_closed()
raise
return self
[docs]
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
if self._listener is not None:
self._listener.close()
await self._listener.wait_closed()
self._listener = None
self._ssh_conn.close()
await self._ssh_conn.wait_closed()