speedtest-ssh 2.6.2__tar.gz → 2.7.1__tar.gz

{speedtest_ssh-2.6.2/speedtest_ssh.egg-info → speedtest_ssh-2.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: speedtest_ssh
-Version: 2.6.2
+Version: 2.7.1
 Summary: A tool meant to check inter-device speeds via ssh
 License: GPL
 Project-URL: Homepage, https://github.com/zwimer/speedtest_ssh

speedtest_ssh-2.7.1/speedtest_ssh/_version.py

@@ -0,0 +1 @@
+__version__: str = "2.7.1"

{speedtest_ssh-2.6.2 → speedtest_ssh-2.7.1}/speedtest_ssh/data_transfer.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 from subprocess import CalledProcessError, PIPE
+from logging import getLogger, DEBUG, INFO
 from typing import TYPE_CHECKING
 from datetime import datetime
 from pathlib import Path
@@ -46,31 +47,29 @@ class DataTransfer:
     This class is a context manager; on exit will remove the remote file on deletion
     """
 
-    def __init__(self, config: Config, verbose: bool):
+    _LOG = "DataTransfer"
+
+    def __init__(self, config: Config):
         """
         :param config: The Config object the DataTransfer instance should use
-        :param verbose: If true, be more verbose
         """
-        self._verbose = verbose
+        self._l = getLogger(self._LOG)
         rand = lambda: random.choice(string.ascii_uppercase + string.ascii_lowercase + string.digits)
         self._remote_f: str = f"/tmp/speedtest-ssh_{datetime.now()}_{''.join(rand() for _ in range(8))}.tmp"
         self._remote_f = self._remote_f.replace(":", "-").replace(" ", "_")
         # We promise that _remote_f components match: ^[a-zA-Z\d_.-]+$ (old rsync args suck)
-        if self._verbose:
-            print("Parsing ssh config and loading keys...")
+        self._l.debug("Parsing ssh config and loading keys...")
         self._sftp_cm = sftp_wrapper(config)
         self._sftp: SFTPClient  # Defined in __enter__
 
     def __enter__(self) -> DataTransfer:
-        if self._verbose:
-            print("Establishing SFTP connection...")
+        self._l.debug("Chosen remote file name: %s", self._remote_f)
         self._sftp = self._sftp_cm.__enter__()
         return self
 
     def __exit__(self, *args, **kwargs):
         self.clean_remote()
-        if self._verbose:
-            print("Terminating SFTP connection...")
+        self._l.debug("Terminating SFTP connection...")
         self._sftp_cm.__exit__(*args, **kwargs)
 
     def put(self, local: Path) -> None:
@@ -90,8 +89,7 @@ class DataTransfer:
         Removes the remote file
         """
         try:
-            if self._verbose:
-                print(f"Removing remote file (if it exists): {self._remote_f}")
+            self._l.debug("Removing remote file (if it exists): %s", self._remote_f)
             self._sftp.remove(self._remote_f)
         except FileNotFoundError:
             pass
@@ -117,22 +115,23 @@ class Rsync(DataTransfer):
     Requires rsync be installed
     """
 
-    def __init__(self, config: Config, verbose: bool):
+    def __init__(self, config: Config):
         """
         :param host: The hostname ssh uses for the remote client
-        :param verbose: If true, be more verbose
         """
-        super().__init__(config, verbose)
-        self._remote_f = re.sub("[^a-zA-Z\\d_-]", "_", self._remote_f)  # Old rsync args suck
+        super().__init__(config)
+        self._remote_f = re.sub(r"[^a-zA-Z\d_/-]", "_", self._remote_f)  # Old rsync args suck
         self._rsync = find_exe("rsync")
         # Determine rsync version
         try:
-            p = run_cmd(self._rsync, "--version", verbose=self._verbose, stdout=PIPE, check=True)
+            p = run_cmd(self._rsync, "--version", stdout=PIPE, check=True)
             output = p.stdout.decode().split("version")[1].split("protocol")[0].strip()
             version: tuple[int, ...] = tuple(int(i) for i in output.split("."))
             self._old = version <= (3, 1, 0)  # macOS has an old version by default
             if self._old:
                 print("\tOld version of rsync detected. Output will be more verbose.")
+            else:
+                self._l.debug("Newer version of rsync detected")
         except (CalledProcessError, KeyError) as e:
             raise RuntimeError(f"Could not determine rsync version from {self._rsync}") from e
         # Determine rsync command
@@ -143,8 +142,10 @@ class Rsync(DataTransfer):
             self._flags = [sshpass, "-e"] + self._flags  # type: ignore
             self._env["SSHPASS"] = config.password
         self._flags.extend(("-hh", "--progress" if self._old else "--info=progress2"))
-        if self._verbose:
-            self._flags.extend(["--verbose", "-vvv"])
+        if self._l.isEnabledFor(INFO):
+            self._flags.append("--verbose")
+        if self._l.isEnabledFor(DEBUG):
+            self._flags.append("-vvv")
         if config.port is not None:
             self._flags.append(f"--port={config.port}")
         # Determine host info
@@ -156,7 +157,7 @@ class Rsync(DataTransfer):
         :param dst: The location to rsync src to
         """
         try:
-            _ = run_cmd(self._rsync, *self._flags, src, dst, verbose=self._verbose, env=self._env, check=True)
+            _ = run_cmd(self._rsync, *self._flags, src, dst, env=self._env, check=True)
         except CalledProcessError as e:
             raise RuntimeError("Failed to transfer data during speed test") from e
 

{speedtest_ssh-2.6.2 → speedtest_ssh-2.7.1}/speedtest_ssh/ping.py

@@ -1,19 +1,23 @@
-from subprocess import CalledProcessError, PIPE
+from subprocess import CalledProcessError
+from logging import INFO
 from sys import platform
 
-from .util import find_exe, run_cmd
+from .util import find_exe, tee_cmd
 
 
 __all__ = ("PingFailed", "ping")
 
 
+_LOG = "Ping"
+
+
 class PingFailed(RuntimeError):
     """
     Raise on ping failure
     """
 
 
-def ping(host: str, verbose: bool, n_pings: int = 10, max_wait: int = 3) -> float:
+def ping(host: str, n_pings: int = 10, max_wait: int = 3) -> float:
     """
     :param host: The host to ping
     :return: The number Ping host
@@ -22,16 +26,13 @@ def ping(host: str, verbose: bool, n_pings: int = 10, max_wait: int = 3) -> floa
     print("Pinging...")
     exe = find_exe("ping")
     try:
-        wait: str = ("-W" if platform == "darwin" else "-w") + str(max_wait)
-        q = [] if verbose else ["-q"]
-        p = run_cmd(exe, "-i.1", f"-c{n_pings}", wait, *q, host, verbose=verbose, stdout=PIPE)
+        wait: str = f"-W{max_wait*1000}" if platform == "darwin" else f"-w{max_wait}"
+        p, stdout = tee_cmd(exe, "-i.1", f"-c{n_pings}", wait, host, level=INFO)
         if p.returncode:
             raise PingFailed(f"{exe} exit code: {p.returncode}")
-        output: str = p.stdout.decode().strip()
+        output: str = stdout.strip()
     except CalledProcessError as e:
         raise PingFailed("Unkown error") from e
-    if verbose:
-        print(output)
     try:
         return float(output.split("/")[-3])
     except IndexError as e:

{speedtest_ssh-2.6.2 → speedtest_ssh-2.7.1}/speedtest_ssh/sftp_wrapper.py

@@ -1,6 +1,7 @@
 from collections.abc import Generator
 from contextlib import contextmanager
 from typing import TYPE_CHECKING
+from logging import getLogger
 from pathlib import Path
 import dataclasses
 
@@ -16,17 +17,22 @@ if TYPE_CHECKING:
 __all__ = ("sftp_wrapper",)
 
 
+_LOG = "SFTPWrapper"
+
+
 def create_paramiko_config(raw: Config) -> dict:
     """
     Create a dict for paramiko to use as its config
     :return: A dict containing info needed for paramiko to ssh
     """
+    log = getLogger(_LOG)
     config: dict[str, str | int | None] = dataclasses.asdict(raw)
     config["username"] = config.pop("user", None)
     config["hostname"] = config.pop("host").lower()  # type: ignore
     # Loads ssh config
     config_f: Path = Path.home() / ".ssh/config"
     if config_f.exists():
+        log.debug("Loading SSH config from: %s", config_f)
         ssh_config: SSHConfig = SSHConfig.from_path(str(config_f))
         global_config: SSHConfigDict = ssh_config.lookup("*")
         host_config: SSHConfigDict = ssh_config.lookup(raw.host)
@@ -36,6 +42,7 @@ def create_paramiko_config(raw: Config) -> dict:
                 value = host_config.get(name, None)
                 value = global_config.get(name, None) if value is None else value
                 if value is not None:
+                    log.debug("  %s: %s", cname, value)
                     config[cname] = value
 
         load("port", "port")
@@ -43,6 +50,7 @@ def create_paramiko_config(raw: Config) -> dict:
         load("key_filename", "identityfile")
     # Install fallbacks
     if config.get("port", None) is None:
+        log.debug("Assuming port 22")
         config["port"] = 22
     return config
 
@@ -51,11 +59,14 @@ def load_keys(client: SSHClient, host: str, port: int) -> None:
     """
     Load available ssh keys into the ssh client
     """
+    log = getLogger(_LOG)
     hosts_f = Path.home() / ".ssh/known_hosts"
     if hosts_f.exists():
+        log.debug("Installing known hosts from: %s", hosts_f)
         client.load_host_keys(str(hosts_f))
     keys: HostKeys = client.get_host_keys()
     # Paramiko seems overly strict about ssh keys in a way ssh is not
+    log.debug("Relaxing paramiko's knownhosts configuration to align with ssh's")
     for i in (host, host.lower()):
         if i in keys:
             keys[f"[{i}]:{port}"] = keys[i]
@@ -66,9 +77,12 @@ def sftp_wrapper(raw: Config) -> Generator["SFTPClient", None, None]:
     """
     A context manager which yields a conncted SFTPClient
     """
+    log = getLogger(_LOG)
+    log.debug("Initializing ssh client")
     with SSHClient() as client:
         config: dict[str, str | int | None] = create_paramiko_config(raw)
         load_keys(client, raw.host, config["port"])  # type: ignore
         client.connect(**config)  # type: ignore
+        log.debug("Initializing SFTP client")
         with client.open_sftp() as sftp:
             yield sftp

{speedtest_ssh-2.6.2 → speedtest_ssh-2.7.1}/speedtest_ssh/speedtest_ssh.py

@@ -1,7 +1,9 @@
+from logging import getLogger, basicConfig, WARNING, DEBUG, INFO
 from tempfile import NamedTemporaryFile as NTF
+from dataclasses import fields
 from pathlib import Path
-from math import log
 import argparse
+import math
 import time
 import sys
 import os
@@ -14,9 +16,10 @@ from ._version import __version__
 
 base_size = 2 * (1024**2)
 _password_env_name = "SPEEDTEST_SSH_PASSWORD"
+_LOG = "speedtest_ssh"
 
 
-def _iteration(temp: Path, client: DataTransfer, size: int, verbose: bool) -> tuple[int, int]:
+def _iteration(temp: Path, client: DataTransfer, size: int) -> tuple[int, int]:
     """
     Time uploading and downloading a file of size bytes
     :param temp: A temporary file we have ownership of
@@ -25,8 +28,8 @@ def _iteration(temp: Path, client: DataTransfer, size: int, verbose: bool) -> tu
     :return: <nanoseconds to upload>, <nanoseconds to download>
     """
     # Create the file to upload; sendfile not always supported so we do this the hard way
-    if verbose:
-        print(f"Filling {temp} with {size} from /dev/urandom... ", end="", flush=True)
+    log = getLogger(_LOG)
+    log.debug("Filling %s with %s MiB from /dev/urandom... ", temp, size / 2**20)
     block: int = min(base_size, 1024**2)
     with temp.open("ab") as f:
         done: int = os.fstat(f.fileno()).st_size  # .tell() can be odd in append mode
@@ -35,18 +38,19 @@ def _iteration(temp: Path, client: DataTransfer, size: int, verbose: bool) -> tu
                 f.write(ur.read(block))
                 done += block
         f.truncate(size)
-    if verbose:
-        print("Done", flush=True)
     # Time uploading the file
+    log.debug("Starting upload test")
     client.clean_remote()
     start: int = time.time_ns()
     client.put(temp)
     up: int = time.time_ns() - start
     # Time downloading the file
+    log.debug("Time taken: %s ns\nStarting download test", up)
     temp.unlink()
     start = time.time_ns()
     client.get(temp)
     down: int = time.time_ns() - start
+    log.debug("Time taken: %s ns\n", down)
     return up, down
 
 
@@ -59,40 +63,29 @@ def _print_results(ping_delay: float | None, up_ns: int, down_ns: int, size: int
     print(f"{prefix}Upload Speed: {fmt(up_ns)}\nDownload Speed: {fmt(down_ns)}")
 
 
-def speedtest_ssh(
-    host: str, num_seconds: int, ping: bool, verbose: bool, mode: str, **kwargs: int | str | None
-) -> None:
+def speedtest_ssh(num_seconds: int, ping: bool, mode: str, conf: Config) -> None:
     """
     Run speedtest_ssh
     """
-    kwargs["password"] = None
-    if kwargs.pop("password_env"):
-        kwargs["password"] = os.environ.get(_password_env_name, None)
-        if kwargs["password"] is None:
-            raise RuntimeError(f"{_password_env_name} is not set")
-    ping_delay: float | None = ping_test(host, verbose) if ping else None
+    log = getLogger(_LOG)
+    ping_delay: float | None = ping_test(conf.host) if ping else None
     print("Connecting...")
-    with (Rsync if mode == "rsync" else SFTP)(
-        Config(host=host, **kwargs),  # type: ignore
-        verbose=verbose,
-    ) as remote:
+    with (Rsync if mode == "rsync" else SFTP)(conf) as remote:
         with NTF(prefix="speedtest_ssh.", dir="/tmp", delete_on_close=False) as ntf:
             ntf.close()
             temp = Path(ntf.name)
-            if verbose:
-                print(f"Created temporary file: {temp}")
+            log.debug("Created temporary file: %s", temp)
             print("Testing...")
             nano_sec: int = 0
             size: int = base_size
             remaining: int = int(1e9) * num_seconds
             while int(1.5 * nano_sec) < remaining - 1:  # Try to get close to seconds
-                size *= 2 ** (0 if not nano_sec else max(1, int(log(remaining / nano_sec, 2))))
+                size *= 2 ** (0 if not nano_sec else max(1, int(math.log(remaining / nano_sec, 2))))
                 # ^ Faster than just *=2
-                up_t, down_t = _iteration(temp, remote, size, verbose)
+                up_t, down_t = _iteration(temp, remote, size)
                 nano_sec = up_t + down_t
                 remaining -= nano_sec
-    if verbose:
-        print(f"\n{'*'*30} Final Results {'*'*30}")
+    log.debug(f"\n{'*'*30} Final Results {'*'*30}")
     _print_results(ping_delay, up_t, down_t, size)
 
 
@@ -106,6 +99,7 @@ def main(argv: list[str]) -> None:
     parser.add_argument("host", help="The host to speedtest the conection to")
     parser.add_argument("--ping", action="store_true", help="Ping test the host as well")
     parser.add_argument("--verbose", action="store_true", help="Be verbose")
+    parser.add_argument("--very-verbose", action="store_true", help="Be very verbose")
     parser.add_argument("-u", "--user", help="The user used for ssh")
     parser.add_argument(
         "--password-env", action="store_true", help=f"Read password from {_password_env_name} environment variable"
@@ -117,7 +111,21 @@ def main(argv: list[str]) -> None:
     parser.add_argument(
         "-m", "--mode", choices=["rsync", "sftp"], default="rsync", help="The speedtest method. Defaults to rsync"
     )
-    return speedtest_ssh(**vars(parser.parse_args(argv[1:])))
+    ns = vars(parser.parse_args(argv[1:]))
+    basicConfig(level=WARNING, format="%(message)s")
+    log = getLogger(_LOG)
+    if ns.pop("verbose"):
+        getLogger().setLevel(INFO)
+    if ns.pop("very_verbose"):
+        getLogger().setLevel(DEBUG)
+        log.debug("Very verbose logging enabled")
+    ns["password"] = None
+    if ns.pop("password_env"):
+        ns["password"] = os.environ.get(_password_env_name, None)
+        if ns["password"] is None:
+            raise RuntimeError(f"{_password_env_name} is not set")
+    conf = Config(**{i.name: ns.pop(i.name) for i in fields(Config)})  # type: ignore
+    return speedtest_ssh(**ns, conf=conf)  # type: ignore
 
 
 def cli() -> None:

speedtest_ssh-2.7.1/speedtest_ssh/util.py

@@ -0,0 +1,66 @@
+from logging import getLogger, DEBUG
+from pathlib import Path
+from shutil import which
+import subprocess
+import shlex
+import os
+
+
+__all__ = ("find_exe", "run_cmd")
+
+
+_LOG = "Util"
+
+
+def find_exe(name: str) -> Path:
+    """
+    :param name: The name of the executable to find
+    :return: The path to the executable
+    """
+    where: str | None = which(name)
+    if where is None:
+        raise RuntimeError(f"Cannot find {name} executable")
+    exe: Path = Path(where).resolve()
+    if not os.access(exe, os.X_OK):
+        raise RuntimeError(f"Cannot find valid {name} executable")
+    return exe
+
+
+def _log_cmd(cmd: Path, args: tuple[str | Path, ...]) -> None:
+    log = getLogger(_LOG)
+    if log.isEnabledFor(DEBUG):
+        full = (cmd, *args)
+        msg = f"{'*'*30} Running Command: {cmd.name} {'*'*30}"
+        msg = f"{msg}\n{' '.join(shlex.quote(str(i)) for i in full)}\n{'*'*len(msg)}"
+        log.debug(msg)
+
+
+def run_cmd(cmd: Path, *args: str | Path, **kwargs) -> subprocess.CompletedProcess:
+    """
+    :param cmd: The command to run
+    :param args: The arguments of cmd
+    :param kwargs: Fowarded to subprocess.Popen
+    :return: Completed process
+    """
+    _log_cmd(cmd, args)
+    return subprocess.run((cmd, *args), **kwargs)
+
+
+def tee_cmd(cmd: Path, *args: str | Path, level: int, **kwargs) -> tuple[subprocess.Popen, str]:
+    """
+    :param cmd: The command to run (and info log the output of)
+    :param args: The arguments of cmd
+    :param level: The log level to use
+    :param kwargs: Fowarded to subprocess.Popen aside from test and stdout
+    :return: Completed process, stdout (do not read stdout from CompletedProcess)
+    """
+    assert "text" not in kwargs
+    assert "stdout" not in kwargs
+    log = getLogger(_LOG)
+    _log_cmd(cmd, args)
+    p = subprocess.Popen((cmd, *args), text=True, stdout=subprocess.PIPE, **kwargs)
+    stdout = ""
+    for line in p.stdout:  # type: ignore
+        log.log(level, line[:-1])
+        stdout += line
+    return p, stdout

{speedtest_ssh-2.6.2 → speedtest_ssh-2.7.1/speedtest_ssh.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: speedtest_ssh
-Version: 2.6.2
+Version: 2.7.1
 Summary: A tool meant to check inter-device speeds via ssh
 License: GPL
 Project-URL: Homepage, https://github.com/zwimer/speedtest_ssh

speedtest_ssh-2.6.2/speedtest_ssh/_version.py

@@ -1 +0,0 @@
-__version__: str = "2.6.2"

speedtest_ssh-2.6.2/speedtest_ssh/util.py

@@ -1,30 +0,0 @@
-from pathlib import Path
-from shutil import which
-import subprocess
-import shlex
-import os
-
-
-__all__ = ("find_exe", "run_cmd")
-
-
-def find_exe(name: str) -> Path:
-    """
-    :param name: The name of the executable to find
-    :return: The path to the executable
-    """
-    where: str | None = which(name)
-    if where is None:
-        raise RuntimeError(f"Cannot find {name} executable")
-    exe: Path = Path(where).resolve()
-    if not os.access(exe, os.X_OK):
-        raise RuntimeError(f"Cannot find valid {name} executable")
-    return exe
-
-
-def run_cmd(cmd: Path, *args: str | Path, verbose: bool, **kwargs) -> subprocess.CompletedProcess:
-    full = (cmd, *args)
-    if verbose:
-        msg = f"{'*'*30} Running Command: {cmd.name} {'*'*30}"
-        print(f"{msg}\n{' '.join(shlex.quote(str(i)) for i in full)}\n{'*'*len(msg)}")
-    return subprocess.run(full, **kwargs)