sshreader 6.2.0__py3-none-any.whl → 7.0.0__py3-none-any.whl

sshreader/__init__.py

@@ -1,6 +1,6 @@
 """A Python Package for parallelizing ssh connections via threading and multiprocessing.
 """
-# Copyright (C) 2015-2025 Jesse Almanrode
+# Copyright (C) 2015-2026 Jesse Almanrode
 #
 #     This program is free software: you can redistribute it and/or modify
 #     it under the terms of the GNU Lesser General Public License as published by

sshreader/customtypes.py

@@ -1,6 +1,6 @@
 """ Typing module for sshreader
 """
-# Copyright (C) 2015-2025 Jesse Almanrode
+# Copyright (C) 2015-2026 Jesse Almanrode
 #
 #     This program is free software: you can redistribute it and/or modify
 #     it under the terms of the GNU Lesser General Public License as published by
@@ -14,23 +14,23 @@
 #
 #     You should have received a copy of the GNU Lesser General Public License
 #     along with this program.  If not, see <http://www.gnu.org/licenses/>.
-from typing import NamedTuple, Optional, Union
+from typing import NamedTuple
 
 
 class Command(NamedTuple):
     cmd: str
-    stdout: Optional[Union[str, bytes]]
-    stderr: Optional[Union[str, bytes]]
+    stdout: str | bytes | None
+    stderr: str | bytes | None
     return_code: int
 
 
 class EnvVars(NamedTuple):
-    username: Optional[str]
-    agent_keys: Optional[list]
-    dsa_key: Optional[str]
-    ecdsa_key: Optional[str]
-    rsa_key: Optional[str]
+    username: str | None
+    agent_keys: list[object] | None
+    dsa_key: str | None
+    ecdsa_key: str | None
+    rsa_key: str | None
 
 
-Timeout = Union[int, float]
+Timeout = int | float
 TimeoutTuple = tuple[Timeout, Timeout]

sshreader/scripts/__init__.py

@@ -1,6 +1,6 @@
 """ Scripts portion of sshreader package
 """
-# Copyright (C) 2015-2025 Jesse Almanrode
+# Copyright (C) 2015-2026 Jesse Almanrode
 #
 #     This program is free software: you can redistribute it and/or modify
 #     it under the terms of the GNU Lesser General Public License as published by

sshreader/scripts/pydsh.py

@@ -1,6 +1,6 @@
 """ A Pythonic implementation of pdsh powered by sshreader
 """
-# Copyright (C) 2015-2025 Jesse Almanrode
+# Copyright (C) 2015-2026 Jesse Almanrode
 #
 #     This program is free software: you can redistribute it and/or modify
 #     it under the terms of the GNU Lesser General Public License as published by
@@ -19,7 +19,8 @@ import logging
 import os
 import sys
 from collections import defaultdict
-from hashlib import md5
+from pathlib import Path
+from typing import Any, Protocol
 
 import click
 from hostlist import collect_hostlist, expand_hostlist
@@ -28,7 +29,7 @@ import sshreader
 
 # GLOBALS
 __author__ = 'Jesse Almanrode'
-__version__ = '3.2.0'
+__version__ = '4.0.0'
 __examples__ = """\b
 Examples:
     pydsh -w host1,host2,host3 "uname -r"
@@ -45,7 +46,24 @@ log.setLevel(logging.WARNING)
 log.propagate = False  # Keeps our messages out of the root logger.
 
 
-def copy_script(scriptfile, job):
+def _job_output(job: sshreader.ServerJob) -> str:
+    result = job.results[0]
+    output = result.stderr if job.status == 255 else result.stdout
+    if output is None:
+        return ''
+    return output.decode() if isinstance(output, bytes) else output
+
+
+class _SFTPUploader(Protocol):
+    def sftp_put(self, srcfile: str, dstfile: str) -> Any: ...
+
+
+class _ScriptJob(Protocol):
+    _conn: _SFTPUploader
+    name: str
+
+
+def copy_script(scriptfile: str, job: _ScriptJob) -> bool:
     """ Copy script to remote host using SFTP connection
 
     :param scriptfile: Path to script file
@@ -56,23 +74,25 @@ def copy_script(scriptfile, job):
     try:
         job._conn.sftp_put(scriptfile, '/tmp/' + script_name)
     except Exception as err:
-        return err
+        log.exception('Failed to copy script %s to %s: %s', scriptfile, getattr(job, 'name', '<unknown>'), err)
+        return False
+    return True
 
 
-def output(thisjob):
+def output(thisjob: sshreader.ServerJob) -> None:
     """ Print output from jobs as they complete in a format that could be piped to dshbak
 
     :param thisjob: <ServerJob> object
     :return: None
     """
-    result = thisjob.results[0] if thisjob.status == 255 else thisjob.results[0].stdout
+    result = _job_output(thisjob)
     if len(result) != 0:
         for line in result.split('\n'):
             sshreader.echo(str(thisjob.name) + ': ' + str(line))
     return None
 
 
-def dshbak(jobresults):
+def dshbak(jobresults: list[sshreader.ServerJob]) -> None:
     """ Output the results of the jobs grouped by hosts
 
     Similar to piping output to `dshbak`
@@ -81,17 +101,14 @@ def dshbak(jobresults):
     :return: None
     """
     for thisjob in jobresults:
-        if thisjob.status == 255:
-            result = thisjob.results[0]
-        else:
-            result = thisjob.results[0].stdout
+        result = _job_output(thisjob)
         if len(result) != 0:
             click.echo(str('-' * 16) + '\n' + str(thisjob.name) + '\n' + str('-' * 16))
             click.echo(result)
     return None
 
 
-def coalesce(jobresults):
+def coalesce(jobresults: list[sshreader.ServerJob]) -> None:
     """ Output the results of jobs coalescing identical output from hosts
 
     Similar to piping output to `dshbak -c`
@@ -99,23 +116,150 @@ def coalesce(jobresults):
     :param jobresults: List of <ServerJob> objects
     :return: None
     """
-    job_hashes = defaultdict(list)
-    output_hashes = dict()
+    output_groups: dict[str, list[str]] = defaultdict(list)
     for job in jobresults:
-        result = job.results[0] if job.status == 255 else job.results[0].stdout
-        md5sum = md5(result.encode()).hexdigest()
-        job_hashes[md5sum].append(job.name)
-        if md5sum not in output_hashes:
-            output_hashes[md5sum] = result
+        output_groups[_job_output(job)].append(job.name)
 
-    for md5sum, stdout in output_hashes.items():
+    for stdout, hosts in output_groups.items():
         if len(stdout) != 0:
-            click.echo(str('-' * 16) + '\n' + collect_hostlist(job_hashes[md5sum]) + '\n' + str('-' * 16))
+            click.echo(str('-' * 16) + '\n' + collect_hostlist(hosts) + '\n' + str('-' * 16))
             click.echo(stdout)
     return None
 
 
-def validate_hostlist(ctx, param, value):
+def setup_script_file(cmd: str, file_flag: bool, confirm_flag: bool) -> tuple[str | list[str], sshreader.Hook | None]:
+    """ Prepare script file for remote execution
+
+    Validates script file, extracts shebang, and creates pre-hook for copying.
+    Prompts for confirmation unless --confirm flag is set.
+
+    :param cmd: Command string or script path
+    :param file_flag: Boolean indicating if cmd is a file path
+    :param confirm_flag: Boolean indicating if confirmation should be skipped
+    :return: Tuple of (cmd_list, prehook) or (cmd, None) if not a file
+    """
+    if not file_flag:
+        return cmd, None
+
+    script_path = Path(cmd)
+    if not script_path.exists() or not script_path.is_file():
+        raise click.BadParameter('Script file not found: ' + str(script_path))
+
+    script_name = script_path.name
+    log.info('Creating copy_script prehook for: ' + script_name)
+    prehook = sshreader.Hook(copy_script, args=[str(script_path)], ssh_established=True)
+
+    content = script_path.read_text(encoding='utf-8', errors='replace')
+    first_line = content.splitlines()[0] if content.splitlines() else ''
+    if not first_line.startswith('#!'):
+        raise click.UsageError('Script must start with #!')
+
+    # Request confirmation unless --confirm flag is set
+    if not confirm_flag:
+        click.echo(f'\nScript: {script_path.resolve()}')
+        click.echo(f'Interpreter: {first_line[2:].strip()}')
+        if not click.confirm('Execute this script on selected hosts?'):
+            raise click.Abort()
+        click.echo()
+
+    cmd_list = [first_line.split('#!').pop().strip() + ' /tmp/' + script_name, 'rm /tmp/' + script_name]
+    return cmd_list, prehook
+
+
+def create_jobs(hostlist: list[str], cmd: str | list[str], username: str, password: str | None,
+                keyfile: str | None, keypass: str | None, port: int, timeout: int, sha2: bool,
+                prehook: sshreader.Hook | None, posthook: sshreader.Hook | None,
+                file_flag: bool) -> list[sshreader.ServerJob]:
+    """ Create ServerJob objects for each host
+
+    :param hostlist: List of hosts from expanded hostlist expression
+    :param cmd: Command(s) to run on each host
+    :param username: SSH username
+    :param password: SSH password
+    :param keyfile: Path to SSH private key
+    :param keypass: Passphrase for private key
+    :param port: Default SSH port
+    :param timeout: SSH timeout in seconds
+    :param sha2: Use SHA2 for RSA keys
+    :param prehook: Pre-execution hook (if any)
+    :param posthook: Post-execution hook (if any)
+    :param file_flag: Boolean indicating if cmd is from a script file
+    :return: List of ServerJob objects
+    """
+    jobs = []
+    for host in hostlist:
+        # Handle port specified in host string
+        if ':' in host:
+            log.info('SSH Port declared in host: ' + host)
+            host, port_str = host.split(':', 1)
+            try:
+                port = int(port_str)
+            except Exception:
+                raise click.BadParameter('Invalid port in host: ' + host) from Exception
+            log.debug((host, port))
+        else:
+            port = port  # Use default port from args
+
+        # Create job with appropriate authentication
+        if keyfile:
+            job = sshreader.ServerJob(
+                host, cmd, username=username, keyfile=keyfile,
+                key_pass=keypass, combine_output=True, rsa_sha2=sha2,
+                timeout=(0.5, timeout)
+            )
+        else:
+            job = sshreader.ServerJob(
+                host, cmd, username=username, password=password,
+                combine_output=True, rsa_sha2=sha2,
+                timeout=(0.5, timeout)
+            )
+
+        job.ssh_port = port
+        if posthook:
+            log.info('Adding posthook to ServerJob for: ' + host)
+            job.post_hook = posthook
+        if file_flag and prehook:
+            log.info('Adding prehook to ServerJob for: ' + host)
+            job.pre_hook = prehook
+        jobs.append(job)
+
+    return jobs
+
+
+def execute_jobs(jobs: list[sshreader.ServerJob], dshbak_flag: bool, coalesce_flag: bool,
+                 redline_flag: bool) -> list[sshreader.ServerJob]:
+    """ Execute jobs and handle output formatting
+
+    :param jobs: List of ServerJob objects
+    :param dshbak_flag: Format output like dshbak
+    :param coalesce_flag: Coalesce similar output
+    :param redline_flag: Run without thread/process overhead
+    :return: List of completed jobs
+    """
+    log.info(f'Sending {len(jobs)} ServerJobs to sshreader module')
+
+    if dshbak_flag or coalesce_flag:
+        # With output aggregation, use progress bar
+        if redline_flag:
+            jobs_finished = sshreader.sshread(jobs, pcount=0, tcount=0, progress_bar=True)
+        else:
+            jobs_finished = sshreader.sshread(jobs, tcount=0, progress_bar=True)
+
+        if coalesce_flag:
+            coalesce(jobs_finished)
+        else:
+            dshbak(jobs_finished)
+    else:
+        # Without aggregation, output as jobs complete
+        if redline_flag:
+            sshreader.sshread(jobs, pcount=0, tcount=0, print_lock=True)
+        else:
+            sshreader.sshread(jobs, tcount=0, print_lock=True)
+
+    return jobs
+
+
+def validate_hostlist(ctx: click.Context | None, param: click.Parameter | str, value: str) -> list[str]:
     """ Callback for click to expand hostlist expressions or error
 
     :param ctx: Click context
@@ -126,56 +270,20 @@ def validate_hostlist(ctx, param, value):
     try:
         return expand_hostlist(value)
     except Exception:
-        raise click.BadOptionUsage(param, 'Invalid hostlist expression') from None
+        option_name = param if isinstance(param, str) else param.name
+        raise click.BadOptionUsage(option_name, 'Invalid hostlist expression') from None
 
 
-@click.command(epilog=__examples__)
-@click.version_option(version=__version__)
-@click.option('--hostlist', '-w', metavar='EXPR', required=True, callback=validate_hostlist,
-              help='Hostlist expression')
-@click.option('--username', '-u', help='Override ssh username')
-@click.option('--keyfile', '-k', type=click.Path(exists=True, dir_okay=False), help='Private key file')
-@click.option('--keypass', '-K', is_flag=True, help='Prompt for private key password')
-@click.option('--prompt', '-p', is_flag=True, help='Prompt for ssh password')
-@click.option('--password', '-P', help='Supply ssh password')
-@click.option('--dshbak', '-D', is_flag=True, help='Group output by host')
-@click.option('--coalesce', '-C', is_flag=True, help='Coalesce similar output from hosts')
-@click.option('--file', '-F', is_flag=True, help='Treat CMD as a script file')
-@click.option('--debug', '-d', is_flag=True, help='Enable debug output')
-@click.option('--verbose', '-v', count=True, help='Increase debug verbosity')
-@click.option('--redline', is_flag=True, help='Run pydsh faster')
-@click.option('--port', default=22, help='SSH Port')
-@click.option('--sha2', is_flag=True, default=True, help='Use SHA2 Hash Algorithm for Keys')
-@click.argument('cmd', nargs=1)
-def cli(**kwargs):
-    """  Run ssh commands in parallel across hosts
-    """
-    if kwargs['debug']:
-        log.setLevel(logging.INFO)
-        if kwargs['verbose']:
-            logging.getLogger('sshreader').setLevel(logging.INFO)
-        if kwargs['verbose'] > 1:
-            log.setLevel(logging.DEBUG)
-        if kwargs['verbose'] > 2:
-            logging.getLogger('sshreader').setLevel(logging.DEBUG)
-    log.debug(kwargs)
-    if kwargs['file']:
-        mkpath = click.Path(exists=True, dir_okay=False)
-        script_path = mkpath(kwargs['cmd'])
-        script_name = os.path.split(script_path)[1]
-        log.info('Creating copy_script prehook for: ' + script_name)
-        prehook = sshreader.Hook(copy_script, args=[script_path], ssh_established=True)
-        with open(script_path) as s:
-            script = s.readline()
-        if script.startswith('#!') is False:
-            raise click.UsageError('Script must start with #!')
-        kwargs['cmd'] = [script.split('#!').pop().strip() + ' /tmp/' + script_name, 'rm /tmp/' + script_name]
-    sshenv = sshreader.envvars()
-    log.debug(sshenv)
+def setup_authentication(kwargs: dict[str, Any], sshenv: sshreader.EnvVars) -> dict[str, Any]:
+    """ Configure SSH authentication credentials
 
-    if kwargs['port'] <= 0 or not isinstance(kwargs['port'], int):
-        raise click.BadOptionUsage('port', 'Please enter a positive integer')
+    Handles the logic for determining which authentication method to use
+    (keyfile, password, SSH agent, or prompt for password).
 
+    :param kwargs: Dictionary of command-line arguments
+    :param sshenv: sshreader.EnvVars object with environment SSH info
+    :return: Updated kwargs with authentication credentials configured
+    """
     if kwargs['username'] is None:
         if sshenv.username is None:
             raise click.ClickException('Unable to determine ssh username. Please provide one using --username')
@@ -196,7 +304,7 @@ def cli(**kwargs):
                 kwargs['keyfile'] = sshenv.dsa_key
                 log.info('Using DSA private key file')
         else:
-            if len(sshenv.agent_keys) == 0:
+            if sshenv.agent_keys is None or len(sshenv.agent_keys) == 0:
                 if not all((kwargs['username'], kwargs['password'])):
                     raise click.ClickException('Unable to find ssh key to use and password not supplied.')
             else:
@@ -214,7 +322,7 @@ def cli(**kwargs):
     # If you specify a password or prompt for one it overrides the ssh key
     if not kwargs['password']:
         if not kwargs['prompt']:
-            if not kwargs['keyfile'] and len(sshenv.agent_keys) == 0:
+            if not kwargs['keyfile'] and (sshenv.agent_keys is None or len(sshenv.agent_keys) == 0):
                 raise click.ClickException('Unable to find ssh key to use and password not supplied or prompt enabled.')
         else:
             log.info('Prompting for password and disabling discovered SSH keyfiles')
@@ -226,46 +334,75 @@ def cli(**kwargs):
         # You provided a password, ignore the SSH key
         kwargs['keyfile'] = None
 
+    return kwargs
+
+
+@click.command(epilog=__examples__)
+@click.version_option(version=__version__)
+@click.option('--hostlist', '-w', metavar='EXPR', required=True, callback=validate_hostlist,
+              help='Hostlist expression')
+@click.option('--username', '-u', help='Override ssh username')
+@click.option('--keyfile', '-k', type=click.Path(exists=True, dir_okay=False), help='Private key file')
+@click.option('--keypass', '-K', is_flag=True, help='Prompt for private key password')
+@click.option('--prompt', '-p', is_flag=True, help='Prompt for ssh password')
+@click.option('--password', '-P', help='Supply ssh password')
+@click.option('--dshbak', '-D', is_flag=True, help='Group output by host')
+@click.option('--coalesce', '-C', is_flag=True, help='Coalesce similar output from hosts')
+@click.option('--file', '-F', is_flag=True, help='Treat CMD as a script file')
+@click.option('--confirm', is_flag=True, help='Skip confirmation prompt when executing scripts')
+@click.option('--debug', '-d', is_flag=True, help='Enable debug output')
+@click.option('--verbose', '-v', count=True, help='Increase debug verbosity')
+@click.option('--redline', is_flag=True, help='Run pydsh faster')
+@click.option('--port', default=22, type=click.IntRange(1, 65535), help='SSH Port')
+@click.option('--timeout', default=30, type=click.IntRange(1, 3600), help='SSH connection timeout in seconds')
+@click.option('--sha2', is_flag=True, default=True, help='Use SHA2 Hash Algorithm for Keys')
+@click.argument('cmd', nargs=1)
+def cli(**kwargs: Any) -> None:
+    """  Run ssh commands in parallel across hosts
+    """
+    # Setup logging
+    if kwargs['debug']:
+        log.setLevel(logging.INFO)
+        if kwargs['verbose']:
+            logging.getLogger('sshreader').setLevel(logging.INFO)
+        if kwargs['verbose'] > 1:
+            log.setLevel(logging.DEBUG)
+        if kwargs['verbose'] > 2:
+            logging.getLogger('sshreader').setLevel(logging.DEBUG)
     log.debug(kwargs)
-    posthook = sshreader.Hook(target=output)
-    jobs = list()
-    for host in kwargs['hostlist']:
-        if ':' in host:
-            log.info('SSH Port declared in host: ' + host)
-            host, port = host.split(':')
-            log.debug((host, port))
-        else:
-            port = kwargs['port']
-        if kwargs['keyfile']:
-            job = sshreader.ServerJob(host, kwargs['cmd'], username=kwargs['username'], keyfile=kwargs['keyfile'],
-                                      key_pass=kwargs['keypass'], combine_output=True, rsa_sha2=kwargs['sha2'])
-        else:
-            job = sshreader.ServerJob(host, kwargs['cmd'], username=kwargs['username'], password=kwargs['password'],
-                                      combine_output=True, rsa_sha2=kwargs['sha2'])
-        job.ssh_port = port
-        if kwargs['dshbak'] is False and kwargs['coalesce'] is False:
-            log.info('Adding posthook to ServerJob for: ' + host)
-            job.post_hook = posthook
-        if kwargs['file']:
-            log.info('Adding prehook to ServerJob for: ' + host)
-            job.pre_hook = prehook
-        jobs.append(job)
 
-    log.info(f'Sending {len(jobs)} ServerJobs to sshreader module')
+    # Setup script file if provided
+    cmd, prehook = setup_script_file(kwargs['cmd'], kwargs['file'], kwargs['confirm'])
+
+    # Configure authentication
+    sshenv = sshreader.envvars()
+    log.debug(sshenv)
+    kwargs = setup_authentication(kwargs, sshenv)
+    log.debug(kwargs)
+
+    # Setup output formatting
+    posthook = None
     if kwargs['dshbak'] is False and kwargs['coalesce'] is False:
-        if kwargs['redline']:
-            sshreader.sshread(jobs, pcount=0, tcount=0, print_lock=True)
-        else:
-            sshreader.sshread(jobs, tcount=0, print_lock=True)
-    else:
-        if kwargs['redline']:
-            jobs_finished = sshreader.sshread(jobs, pcount=0, tcount=0, progress_bar=True)
-        else:
-            jobs_finished = sshreader.sshread(jobs, tcount=0, progress_bar=True)
-        if kwargs['coalesce']:
-            coalesce(jobs_finished)
-        else:
-            dshbak(jobs_finished)
+        posthook = sshreader.Hook(target=output)
+
+    # Create jobs for all hosts
+    jobs = create_jobs(
+        hostlist=kwargs['hostlist'],
+        cmd=cmd,
+        username=kwargs['username'],
+        password=kwargs['password'],
+        keyfile=kwargs['keyfile'],
+        keypass=kwargs['keypass'],
+        port=kwargs['port'],
+        timeout=kwargs['timeout'],
+        sha2=kwargs['sha2'],
+        prehook=prehook,
+        posthook=posthook,
+        file_flag=kwargs['file']
+    )
+
+    # Execute jobs
+    execute_jobs(jobs, kwargs['dshbak'], kwargs['coalesce'], kwargs['redline'])
     sys.exit(0)