diff --git a/python/binharness/agentenvironment.py b/python/binharness/agentenvironment.py index 8b24498..4f6e4ee 100644 --- a/python/binharness/agentenvironment.py +++ b/python/binharness/agentenvironment.py @@ -115,22 +115,31 @@ def __init__( # noqa: PLR0913 self._pid = pid @cached_property - def stdin(self: AgentProcess) -> AgentIO: + def stdin(self: AgentProcess) -> AgentIO | None: """Get the standard input stream of the process.""" - fd = self._client.get_process_channel(self._env_id, self._pid, 0) - return AgentIO(self._client, self._env_id, fd) + try: + fd = self._client.get_process_channel(self._env_id, self._pid, 0) + return AgentIO(self._client, self._env_id, fd) + except RuntimeError: + return None # TODO: verify that this is the right error @cached_property - def stdout(self: AgentProcess) -> AgentIO: + def stdout(self: AgentProcess) -> AgentIO | None: """Get the standard output stream of the process.""" - fd = self._client.get_process_channel(self._env_id, self._pid, 1) - return AgentIO(self._client, self._env_id, fd) + try: + fd = self._client.get_process_channel(self._env_id, self._pid, 1) + return AgentIO(self._client, self._env_id, fd) + except RuntimeError: + return None # TODO: verify that this is the right error @cached_property - def stderr(self: AgentProcess) -> AgentIO: + def stderr(self: AgentProcess) -> AgentIO | None: """Get the standard error stream of the process.""" - fd = self._client.get_process_channel(self._env_id, self._pid, 2) - return AgentIO(self._client, self._env_id, fd) + try: + fd = self._client.get_process_channel(self._env_id, self._pid, 2) + return AgentIO(self._client, self._env_id, fd) + except RuntimeError: + return None # TODO: verify that this is the right error @property def returncode(self: AgentProcess) -> int | None: @@ -174,7 +183,7 @@ def run_command( stdin=True, stdout=True, stderr=True, - executable=str(args[0]), + executable=str(normalized_args[0]), env=list(env.items()) if env else None, cwd=str(cwd) if cwd else None, setuid=None, diff --git a/python/binharness/common/busybox.py b/python/binharness/common/busybox.py index a013fcb..f4193fa 100644 --- a/python/binharness/common/busybox.py +++ b/python/binharness/common/busybox.py @@ -4,7 +4,7 @@ from os import environ from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from binharness.types import InjectableExecutor, Target from binharness.types.injection import ExecutableInjection @@ -48,7 +48,8 @@ def mktemp( """Run mktemp in the environment and returns the Path created.""" proc = self.run("mktemp", "-d") if directory else self.run("mktemp") stdout, _ = proc.communicate() - return Path(stdout.decode().strip()) + # TODO: Find out how to not cast + return Path(cast(bytes, stdout).decode().strip()) def shell( self: BusyboxInjection, command: str, env: dict[str, str] | None = None diff --git a/python/binharness/localenvironment.py b/python/binharness/localenvironment.py index a4f822a..988034f 100644 --- a/python/binharness/localenvironment.py +++ b/python/binharness/localenvironment.py @@ -110,24 +110,18 @@ def __init__( ) @property - def stdin(self: LocalProcess) -> IO[bytes]: + def stdin(self: LocalProcess) -> IO[bytes] | None: """Get the standard input stream of the process.""" - if self.popen.stdin is None: - raise ValueError # pragma: no cover return self.popen.stdin @property - def stdout(self: LocalProcess) -> IO[bytes]: + def stdout(self: LocalProcess) -> IO[bytes] | None: """Get the standard output stream of the process.""" - if self.popen.stdout is None: - raise ValueError # pragma: no cover return self.popen.stdout @property - def stderr(self: LocalProcess) -> IO[bytes]: + def stderr(self: LocalProcess) -> IO[bytes] | None: """Get the standard error stream of the process.""" - if self.popen.stderr is None: - raise ValueError # pragma: no cover return self.popen.stderr @property diff --git a/python/binharness/types/process.py b/python/binharness/types/process.py index 168b3bb..d676638 100644 --- a/python/binharness/types/process.py +++ b/python/binharness/types/process.py @@ -34,17 +34,17 @@ def __init__( self.cwd = cwd or environment.get_tempdir() @abstractproperty - def stdin(self: Process) -> IO[bytes]: + def stdin(self: Process) -> IO[bytes] | None: """Get the standard input stream of the process.""" raise NotImplementedError @abstractproperty - def stdout(self: Process) -> IO[bytes]: + def stdout(self: Process) -> IO[bytes] | None: """Get the standard output stream of the process.""" raise NotImplementedError @abstractproperty - def stderr(self: Process) -> IO[bytes]: + def stderr(self: Process) -> IO[bytes] | None: """Get the standard error stream of the process.""" raise NotImplementedError @@ -65,10 +65,13 @@ def wait(self: Process, timeout: float | None = None) -> int: def communicate( self: Process, input_: bytes | None = None, timeout: float | None = None - ) -> tuple[bytes, bytes]: + ) -> tuple[bytes | None, bytes | None]: """Send input to the process and return its output and error streams.""" - if input_ is not None: - self.stdin.write(input_) - self.stdin.close() + if self.stdin is not None: + if input_ is not None: + self.stdin.write(input_) + self.stdin.close() self.wait(timeout) - return (self.stdout.read(), self.stderr.read()) + stdout = self.stdout.read() if self.stdout is not None else None + stderr = self.stderr.read() if self.stderr is not None else None + return (stdout, stderr) diff --git a/python/tests/bootstrap/test_docker.py b/python/tests/bootstrap/test_docker.py index 062952d..d0f2559 100644 --- a/python/tests/bootstrap/test_docker.py +++ b/python/tests/bootstrap/test_docker.py @@ -27,6 +27,7 @@ def test_bootstrap_env_from_image( proc.wait() assert proc.returncode == 0 + assert proc.stdout is not None assert proc.stdout.read() == b"hello world\n" finally: agent.container.kill() diff --git a/python/tests/bootstrap/test_ssh.py b/python/tests/bootstrap/test_ssh.py index 90aa186..06487c3 100644 --- a/python/tests/bootstrap/test_ssh.py +++ b/python/tests/bootstrap/test_ssh.py @@ -47,6 +47,7 @@ def test_bootstrap_ssh_environment_with_client( proc.wait() assert proc.returncode == 0 + assert proc.stdout is not None assert proc.stdout.read() == b"hello world\n" finally: agent.stop() diff --git a/python/tests/test_busybox.py b/python/tests/test_busybox.py index d0e55e9..58de298 100644 --- a/python/tests/test_busybox.py +++ b/python/tests/test_busybox.py @@ -43,10 +43,14 @@ def test_nc_interaction() -> None: server_proc = busybox.nc("localhost", 10001, listen=True) client_proc = busybox.nc("localhost", 10001) + assert client_proc.stdin is not None + assert server_proc.stdin is not None client_proc.stdin.write(b"hello\n") client_proc.stdin.flush() server_proc.stdin.write(b"hello back\n") server_proc.stdin.flush() + assert server_proc.stdout is not None assert server_proc.stdout.readline() == b"hello\n" + assert client_proc.stdout is not None assert client_proc.stdout.readline() == b"hello back\n" diff --git a/python/tests/test_localenvironment.py b/python/tests/test_localenvironment.py index 00d7c70..fc968a0 100644 --- a/python/tests/test_localenvironment.py +++ b/python/tests/test_localenvironment.py @@ -43,6 +43,7 @@ def test_stdout() -> None: busybox = BusyboxInjection() busybox.install(env) proc = busybox.shell("echo hello") + assert proc.stdout is not None assert proc.stdout.read() == b"hello\n" @@ -52,6 +53,7 @@ def test_stderr() -> None: busybox = BusyboxInjection() busybox.install(env) proc = busybox.shell("echo hello 1>&2") + assert proc.stderr is not None assert proc.stderr.read() == b"hello\n" @@ -62,8 +64,10 @@ def test_process_poll() -> None: busybox.install(env) proc = busybox.run("head") assert proc.poll() is None + assert proc.stdin is not None proc.stdin.write(b"hello\n") proc.stdin.close() proc.wait() assert proc.poll() is not None + assert proc.stdout is not None assert proc.stdout.read() == b"hello\n"