46 lines
1.2 KiB
Python
46 lines
1.2 KiB
Python
"""SSH and test setup utilities"""
|
|
|
|
from pathlib import Path
|
|
from typing import NamedTuple
|
|
|
|
from .port import find_free_port, setup_port_forwarding
|
|
|
|
|
|
class SSHConnection(NamedTuple):
|
|
host_port: int
|
|
ssh_key: str
|
|
|
|
|
|
def setup_ssh_connection(
|
|
target,
|
|
temp_dir: str,
|
|
assets_ssh_privkey: str,
|
|
) -> SSHConnection:
|
|
"""Set up SSH connection with port forwarding to test VM
|
|
|
|
Args:
|
|
target: Test VM target
|
|
temp_dir: Temporary directory for SSH key
|
|
assets_ssh_privkey: Path to SSH private key asset
|
|
|
|
Returns:
|
|
SSHConnection with host_port and ssh_key path
|
|
"""
|
|
host_port = find_free_port()
|
|
target.wait_for_open_port(22)
|
|
|
|
# Check if this is a container test (doesn't have forward_port method)
|
|
if hasattr(target, "forward_port"):
|
|
setup_port_forwarding(target, host_port)
|
|
else:
|
|
# In container tests, we use the host network namespace
|
|
# so we can connect directly to the container's SSH port
|
|
host_port = 22
|
|
|
|
ssh_key = Path(temp_dir) / "id_ed25519"
|
|
with ssh_key.open("w") as f, Path(assets_ssh_privkey).open() as src:
|
|
f.write(src.read())
|
|
ssh_key.chmod(0o600)
|
|
|
|
return SSHConnection(host_port, str(ssh_key))
|