From 081324aaed14047fb7e6db1b21f52600dbf03948 Mon Sep 17 00:00:00 2001 From: Kevin Phoenix Date: Fri, 6 Oct 2023 13:13:11 -0700 Subject: [PATCH] Remove SSH Environment (#19) --- binharness/__init__.py | 4 +- binharness/environment/__init__.py | 7 - binharness/environment/sshenvironment.py | 232 ------------------ .../{environment => }/localenvironment.py | 2 +- binharness/tests/test_busybox.py | 2 +- binharness/tests/test_executor.py | 2 +- binharness/tests/test_inject.py | 2 +- binharness/tests/test_localenvironment.py | 2 +- binharness/tests/test_sshenvironment.py | 145 ----------- binharness/tests/test_target_serialization.py | 2 +- 10 files changed, 8 insertions(+), 392 deletions(-) delete mode 100644 binharness/environment/__init__.py delete mode 100644 binharness/environment/sshenvironment.py rename binharness/{environment => }/localenvironment.py (98%) delete mode 100644 binharness/tests/test_sshenvironment.py diff --git a/binharness/__init__.py b/binharness/__init__.py index ea85184..45cf2d7 100644 --- a/binharness/__init__.py +++ b/binharness/__init__.py @@ -1,10 +1,10 @@ -"""binharness - A library for analyzing the a program in its environment.""" +"""binharness - A library for analyzing a program in its environment.""" from __future__ import annotations __version__ = "0.1.0dev0" from binharness.common import BusyboxInjection -from binharness.environment import LocalEnvironment +from binharness.localenvironment import LocalEnvironment from binharness.serialize import TargetImportError, export_target, import_target from binharness.types import ( IO, diff --git a/binharness/environment/__init__.py b/binharness/environment/__init__.py deleted file mode 100644 index 888f7cd..0000000 --- a/binharness/environment/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""binharness.environment - Environment implementations for Binharness.""" -from __future__ import annotations - -from .localenvironment import LocalEnvironment -from .sshenvironment import SSHEnvironment - -__all__ = ["LocalEnvironment", "SSHEnvironment"] diff --git a/binharness/environment/sshenvironment.py b/binharness/environment/sshenvironment.py deleted file mode 100644 index fb71746..0000000 --- a/binharness/environment/sshenvironment.py +++ /dev/null @@ -1,232 +0,0 @@ -"""binharness.environment.sshenvironment: SSHEnvironment.""" -from __future__ import annotations - -import shlex -import stat -from pathlib import Path -from typing import TYPE_CHECKING, Any, Sequence - -import paramiko - -from binharness.types.environment import Environment -from binharness.types.process import Process -from binharness.util import join_normalized_args, normalize_args - -if TYPE_CHECKING: - from binharness.types.io import IO - - -class SSHEnvironmentError(Exception): - """An error occurred in an SSHEnvironment.""" - - -class SSHInvalidArgumentsError(SSHEnvironmentError): - """Invalid arguments were passed to an SSHEnvironment.""" - - -class SSHNotConnectedError(SSHEnvironmentError): - """An SSHEnvironment is not connected to a host.""" - - -class SSHPermissionError(SSHEnvironmentError): - """An SSHEnvironment does not have permission to perform an action.""" - - -class SSHProcess(Process): - """A process running in an SSHEnvironment.""" - - _channel: paramiko.Channel - _returncode: int | None - - def __init__( # noqa: PLR0913 - self: SSHProcess, - channel: paramiko.Channel, - environment: SSHEnvironment, - args: Sequence[str], - env: dict[str, str] | None = None, - cwd: Path | None = None, - ) -> None: - """Create a Process.""" - super().__init__(environment, args, env=env, cwd=cwd) - self._channel = channel - self._returncode = None - self._channel.update_environment( - {key.encode(): value.encode() for key, value in self.env.items()} - ) - command_str = join_normalized_args(self.args) - command_str = "cd " + shlex.quote(str(self.cwd)) + " && " + command_str - self._channel.exec_command(command_str) - - @property - def stdin(self: SSHProcess) -> IO[bytes]: - """Get the standard input stream of the process.""" - return self._channel.makefile_stdin("wb") - - @property - def stdout(self: SSHProcess) -> IO[bytes]: - """Get the standard output stream of the process.""" - return self._channel.makefile("rb") - - @property - def stderr(self: SSHProcess) -> IO[bytes]: - """Get the standard error stream of the process.""" - return self._channel.makefile_stderr("rb") - - @property - def returncode(self: SSHProcess) -> int | None: - """Get the process' exit code.""" - return self._returncode - - def poll(self: SSHProcess) -> int | None: - """Return the process' exit code if it has terminated, or None.""" - rc = ( - self._channel.recv_exit_status() - if self._channel.exit_status_ready() - else None - ) - if rc is not None: - self._returncode = rc - return rc - - def wait(self: SSHProcess, timeout: float | None = None) -> int: - """Wait for the process to terminate and return its exit code.""" - if timeout is not None: - self._channel.settimeout(timeout) - while True: - rc = self._channel.recv_exit_status() - if rc is not None: - self._returncode = rc - break - return self._channel.recv_exit_status() - - -class SSHEnvironment(Environment): - """An environment over SSH.""" - - _client: paramiko.SSHClient - _tmp_base: Path - - def __init__( - self: SSHEnvironment, - *, - client: paramiko.SSHClient | None = None, - hostname: str | None = None, - connection_args: dict[str, Any] | None = None, - tmp_base: Path = Path("/tmp"), - ) -> None: - """Create a SSHEnvironment.""" - super().__init__() - if (client is None and hostname is None) or ( - client is not None and hostname is not None - ): - raise SSHInvalidArgumentsError - if client is not None: - self._client = client - if hostname is not None: - self._connect(hostname, connection_args if connection_args else {}) - self._tmp_base = tmp_base - - # API Methods - - def run_command( - self: SSHEnvironment, - *args: Path | str | Sequence[Path | str], - env: dict[str, str] | None = None, - cwd: Path | None = None, - ) -> SSHProcess: - """Run a command in the environment. - - The command is run as a process in the environment. For now, arguments - are passed to Process as-is. - """ - transport = self._client.get_transport() - if transport is None: - raise SSHNotConnectedError - return SSHProcess( - transport.open_session(), self, normalize_args(*args), env=env, cwd=cwd - ) - - def inject_files( - self: SSHEnvironment, - files: list[tuple[Path, Path]], - ) -> None: - """Inject files into the environment. - - The first element of the tuple is the path to the file on the host - machine, the second element is the path to the file in the environment. - """ - sftp = self._client.open_sftp() - for file in files: - dest_path = file[1] - if _is_dir(sftp, dest_path): - dest_path = dest_path / file[0].name - else: - _make_dirs(sftp, dest_path.parent) - sftp.put(str(file[0]), str(dest_path)) - sftp.chmod(str(dest_path), stat.S_IMODE(file[0].stat().st_mode)) - sftp.close() - - def retrieve_files( - self: SSHEnvironment, - files: list[tuple[Path, Path]], - ) -> None: - """Retrieve files from the environment. - - The first element of the tuple is the path to the file in the - environment, the second element is the path to the file on the host - machine. - """ - sftp = self._client.open_sftp() - for file in files: - sftp.get(str(file[0]), str(file[1])) - sftp.close() - - def get_tempdir(self: SSHEnvironment) -> Path: - """Get a Path for a temporary directory.""" - return self._tmp_base - - # Connection Methods - def _connect( - self: SSHEnvironment, host: str, connection_args: dict[str, Any] - ) -> None: # pragma: no cover - """Connect to the environment.""" - self._client = paramiko.SSHClient() - self._client.load_system_host_keys() - self._client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) - self._client.connect(host, **connection_args) - - def close(self: SSHEnvironment) -> None: - """Close the connection to the environment.""" - if hasattr(self, "_client"): - self._client.close() - - def __del__(self: SSHEnvironment) -> None: - """Close the connection to the environment.""" - self.close() - - -# Utility functions for SFTP - - -def _is_dir(sftp: paramiko.SFTPClient, path: Path) -> bool: - try: - file_mode = sftp.stat(str(path)).st_mode - except FileNotFoundError: - return False - if file_mode is None: - raise SSHEnvironmentError - return bool(stat.S_ISDIR(file_mode)) - - -def _make_dirs(sftp: paramiko.SFTPClient, path: Path) -> None: - working_path = path - while len(path.parts) > 1: - try: - sftp.mkdir(str(working_path)) - except PermissionError as err: # noqa: PERF203 - # No permission to create directory - raise SSHPermissionError from err - except FileNotFoundError: # Parent directory does not exist - _make_dirs(sftp, working_path.parent) - except OSError: # Directory already exists - return diff --git a/binharness/environment/localenvironment.py b/binharness/localenvironment.py similarity index 98% rename from binharness/environment/localenvironment.py rename to binharness/localenvironment.py index 45abea4..23b08fb 100644 --- a/binharness/environment/localenvironment.py +++ b/binharness/localenvironment.py @@ -1,4 +1,4 @@ -"""binharness.environment.localenvironment - A local environment.""" +"""binharness.localenvironment - A local environment.""" from __future__ import annotations import shutil diff --git a/binharness/tests/test_busybox.py b/binharness/tests/test_busybox.py index e6db11d..a54ad99 100644 --- a/binharness/tests/test_busybox.py +++ b/binharness/tests/test_busybox.py @@ -1,7 +1,7 @@ from __future__ import annotations from binharness.common.busybox import BusyboxInjection -from binharness.environment.localenvironment import LocalEnvironment +from binharness.localenvironment import LocalEnvironment def test_busybox_injection() -> None: diff --git a/binharness/tests/test_executor.py b/binharness/tests/test_executor.py index 170386b..7da0ef9 100644 --- a/binharness/tests/test_executor.py +++ b/binharness/tests/test_executor.py @@ -5,7 +5,7 @@ import pytest from binharness.common.busybox import BusyboxShellExecutor -from binharness.environment.localenvironment import LocalEnvironment +from binharness.localenvironment import LocalEnvironment from binharness.types.executor import ( ExecutorEnvironmentMismatchError, ) diff --git a/binharness/tests/test_inject.py b/binharness/tests/test_inject.py index 8450ce0..bace783 100644 --- a/binharness/tests/test_inject.py +++ b/binharness/tests/test_inject.py @@ -4,7 +4,7 @@ import pytest -from binharness.environment.localenvironment import LocalEnvironment +from binharness.localenvironment import LocalEnvironment from binharness.types.injection import ( ExecutableInjection, Injection, diff --git a/binharness/tests/test_localenvironment.py b/binharness/tests/test_localenvironment.py index 226a775..821f228 100644 --- a/binharness/tests/test_localenvironment.py +++ b/binharness/tests/test_localenvironment.py @@ -4,7 +4,7 @@ import tempfile from binharness.common.busybox import BusyboxInjection -from binharness.environment.localenvironment import LocalEnvironment +from binharness.localenvironment import LocalEnvironment def test_run_command() -> None: diff --git a/binharness/tests/test_sshenvironment.py b/binharness/tests/test_sshenvironment.py deleted file mode 100644 index 83ee7ed..0000000 --- a/binharness/tests/test_sshenvironment.py +++ /dev/null @@ -1,145 +0,0 @@ -from __future__ import annotations - -import pathlib -import tempfile -from typing import TYPE_CHECKING, Generator - -import mockssh -import pytest - -from binharness.common.busybox import BusyboxInjection -from binharness.environment.sshenvironment import ( - SSHEnvironment, - SSHInvalidArgumentsError, - SSHNotConnectedError, - _is_dir, - _make_dirs, -) - -if TYPE_CHECKING: - import paramiko - - -USER_KEYS = { - "test": str((pathlib.Path(__file__).parent / "ssh_keys" / "test").absolute()), -} - - -@pytest.fixture() -def ssh_server() -> Generator[mockssh.Server, None, None]: - with mockssh.Server(USER_KEYS) as server: - yield server - - -@pytest.fixture() -def ssh_client(ssh_server: mockssh.Server) -> Generator[paramiko.SSHClient, None, None]: - client = ssh_server.client("test") - yield client - client.close() - - -@pytest.fixture() -def sftp_client( - ssh_client: paramiko.SSHClient, -) -> Generator[paramiko.SFTPClient, None, None]: - sftp_client = ssh_client.open_sftp() - yield sftp_client - sftp_client.close() - - -@pytest.fixture() -def env(ssh_server: mockssh.Server) -> Generator[SSHEnvironment, None, None]: - env = SSHEnvironment(client=ssh_server.client("test")) - yield env - env.close() - - -@pytest.fixture() -def busybox(env: SSHEnvironment) -> BusyboxInjection: - busybox = BusyboxInjection() - busybox.install(env) - return busybox - - -def test_run_command(env: SSHEnvironment) -> None: - proc = env.run_command(["echo", "hello"]) - stdout, _ = proc.communicate(timeout=5) - assert proc.returncode == 0 - assert stdout == b"hello\n" - - -def test_inject_files(env: SSHEnvironment) -> None: - env_temp = env.get_tempdir() - with tempfile.TemporaryDirectory() as tmp_dir: - tmp_path = pathlib.Path(tmp_dir) - file = tmp_path / "test.txt" - file.write_text("hello") - env.inject_files([(file, env_temp / "test.txt")]) - - with tempfile.TemporaryDirectory() as tmp_dir: - local_file = pathlib.Path(tmp_dir) / "test.txt" - env.retrieve_files([(env_temp / "test.txt", local_file)]) - assert local_file.read_text() == "hello" - - -def test_get_tempdir(env: SSHEnvironment) -> None: - assert env.get_tempdir() == pathlib.Path("/tmp") - - -def test_stdout(env: SSHEnvironment) -> None: - busybox = BusyboxInjection() - busybox.install(env) - proc = busybox.shell("echo hello") - assert proc.stdout.read() == b"hello\n" - - -def test_stderr(env: SSHEnvironment) -> None: - busybox = BusyboxInjection() - busybox.install(env) - proc = busybox.shell("echo hello 1>&2") - assert proc.stderr.read() == b"hello\n" - - -def test_process_poll(env: SSHEnvironment) -> None: - busybox = BusyboxInjection() - busybox.install(env) - proc = busybox.run("sleep", "0.1") - assert proc.poll() is None - proc.wait() - assert proc.poll() is not None - - -def test_invalid_args() -> None: - with pytest.raises(SSHInvalidArgumentsError): - SSHEnvironment() - - -def test_command_after_close(env: SSHEnvironment) -> None: - env.close() - with pytest.raises(SSHNotConnectedError): - env.run_command(["echo", "hello"]) - - -def test_connect_hostname(ssh_server: mockssh.Server) -> None: - SSHEnvironment( - hostname=ssh_server.host, - connection_args={ - "port": ssh_server.port, - "username": "test", - "key_filename": USER_KEYS["test"], - "allow_agent": False, - "look_for_keys": False, - }, - ).close() - - -def test_is_dir(sftp_client: paramiko.SFTPClient) -> None: - assert _is_dir(sftp_client, pathlib.Path("/tmp")) - - -def test_make_dir(sftp_client: paramiko.SFTPClient) -> None: - _make_dirs(sftp_client, pathlib.Path("/tmp/test")) - - -def test_make_dir_already_exists(sftp_client: paramiko.SFTPClient) -> None: - _make_dirs(sftp_client, pathlib.Path("/")) diff --git a/binharness/tests/test_target_serialization.py b/binharness/tests/test_target_serialization.py index c648350..77936aa 100644 --- a/binharness/tests/test_target_serialization.py +++ b/binharness/tests/test_target_serialization.py @@ -6,7 +6,7 @@ import pytest -from binharness.environment.localenvironment import LocalEnvironment +from binharness.localenvironment import LocalEnvironment from binharness.serialize import TargetImportError, export_target, import_target from binharness.types.executor import NullExecutor from binharness.types.target import Target