From 5b87a67f5b64fe7e002dfd4950a67b70b5fdc0f3 Mon Sep 17 00:00:00 2001 From: Jacob Callahan Date: Tue, 1 Oct 2024 14:36:44 -0400 Subject: [PATCH] Add explicit close and context manager to Connection This commit introduces a new close method that will close the session on a Connection object. However, ssh2-rs states that this will not close the underlying socket. If that's an issue in the future, perhaps we also track the TCP connection and close that along with the session. Since we're now allowing explicit closing, I also added a context manager to the Connection class, so people can handle that implicitly. Fixes #15 --- Cargo.toml | 6 +++++- README.md | 25 ++++++++++++++++++++----- benchmarks/README.md | 2 +- pyproject.toml | 1 + src/connection.rs | 26 ++++++++++++++++++++++++++ tests/test_connection.py | 8 ++++++++ 6 files changed, 61 insertions(+), 7 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2b5585d..fe5fe71 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "hussh" -version = "0.1.7" +version = "0.1.8" edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html @@ -14,3 +14,7 @@ pyo3 = "0.22" # ssh2 = "0.9" # temporary until ssh2#312 makes it into a release. probably 0.9.5 ssh2 = { git = "https://github.com/alexcrichton/ssh2-rs", branch = "master" } + +[profile.release] +lto = true + diff --git a/README.md b/README.md index 3d581de..94ec09c 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,21 @@ Hussh can also do agent-based authentication, if you've already established it. conn = Connection("my.test.server") ``` +## Cleaning up after yourself + +Hussh will clean up after itself automatically when the `Connection` object is garbage collected. + +However, if you want to more explicitly clean up after yourself, you can `close` the connection. +```python +conn.close() +``` +or you can use the `Connection` class' context manager, which will `close` when you exit the context. +```python +with Connection(host="my.test.server", password="pass") as conn: + result = conn.execute("ls") +assert result.status == 0 +``` + # Executing commands The most basic foundation of ssh libraries is the ability to execute commands against the remote host. For Hussh, just use the `Connection` object's `execute` method. @@ -82,7 +97,7 @@ Each execute returns an `SSHResult` object with command's stdout, stderr, and st # SFTP If you need to transfer files to/from the remote host, SFTP may be your best bet. -## Writing Files and Data +## Writing files and data ```python # write a local file to the remote destination conn.sftp_write(local_path="/path/to/my/file", remote_path="/dest/path/file") @@ -91,7 +106,7 @@ conn.sftp_write(local_path="/path/to/my/file", remote_path="/dest/path/file") conn.sftp_write_data(data="Hello there!", remote_path="/dest/path/file") ``` -## Reading Files +## Reading files ```python # You can copy a remote file to a local destination conn.sftp_read(remote_path="/dest/path/file", local_path="/path/to/my/file") @@ -113,7 +128,7 @@ By default, if you don't pass in an alternate `dest_path`, Hussh will copy it to # SCP For remote servers that support SCP, Hussh can do that to. -## Writing Files and Data +## Writing files and data ```python # write a local file to the remote destination conn.scp_write(local_path="/path/to/my/file", remote_path="/dest/path/file") @@ -122,7 +137,7 @@ conn.scp_write(local_path="/path/to/my/file", remote_path="/dest/path/file") conn.scp_write_data(data="Hello there!", remote_path="/dest/path/file") ``` -## Reading Files +## Reading files ```python # You can copy a remote file to a local destination conn.scp_read(remote_path="/dest/path/file", local_path="/path/to/my/file") @@ -157,7 +172,7 @@ print(shell.result.stdout) # Disclaimer This is a VERY early project that should not be used in production code! -There isn't even proper exception handling, so try/except won't work. +There isn't even proper exception handling, so expect some Rust panics to fall through. With that said, try it out and let me know your thoughts! # Future Features diff --git a/benchmarks/README.md b/benchmarks/README.md index 4c36093..8af0a98 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -25,6 +25,6 @@ This will ultimately collect all the benchmark and memray information into a tab Alternatively, if you'd prefer to run individual benchmarks, you can do that. ```bash -python test_hussh.py +python bench_hussh.py ``` This will also create a memray output file for each script ran. diff --git a/pyproject.toml b/pyproject.toml index b9379b7..652ebc0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ classifiers = [ "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Topic :: System :: Shells", "Topic :: System :: Networking", "Topic :: Software Development :: Libraries", diff --git a/src/connection.rs b/src/connection.rs index c9d3dd5..f3c22ac 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -495,6 +495,32 @@ impl Connection { FileTailer::new(self, remote_file, None) } + /// Close the connection's session + fn close(&self) -> PyResult<()> { + self.session + .disconnect(None, "Bye from Hussh", None) + .unwrap(); + Ok(()) + } + + /// Provide an enter for the context manager + fn __enter__(slf: PyRef) -> PyRef { + slf + } + + /// Provide an exit for the context manager + /// This will close the session + #[pyo3(signature = (_exc_type=None, _exc_value=None, _traceback=None))] + fn __exit__( + &mut self, + _exc_type: Option<&Bound<'_, PyAny>>, + _exc_value: Option<&Bound<'_, PyAny>>, + _traceback: Option<&Bound<'_, PyAny>>, + ) -> PyResult<()> { + let _ = self.close(); + Ok(()) + } + fn __repr__(&self) -> PyResult { Ok(format!( "Connection(host={}, port={}, username={}, password=*****)", diff --git a/tests/test_connection.py b/tests/test_connection.py index 6e56141..1c67e3f 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -57,6 +57,14 @@ def test_bad_command(conn): assert "command not found" in result.stderr +def test_conn_context(): + """Test that the Connection class' context manager works.""" + with Connection(host="localhost", port=8022, password="toor") as conn: + result = conn.execute("echo hello") + assert result.status == 0 + assert result.stdout == "hello\n" + + def test_text_scp(conn): """Test that we can copy a file to the server and read it back.""" # copy a local file to the server