Added support for identity files in the CLI and TUI.

Added support for forward ports.

Excplicitly declared password and start commands to be NotYetImplemented.

Author: cornyhorse

cli/sshaman_cli.py

@@ -49,11 +49,11 @@ def initialize_sample(config_path):
 @click.argument('alias')
 @click.argument('host')
 @click.argument('user')
-@click.option('--port', required=False, default=22)
-@click.option('--identity_file', required=False, default='')
-@click.option('--password', required=False, default='')
-@click.option('--forward_ports', required=False)
-@click.option('--start_commands', required=False)
+@click.option('--port', '-p', required=False, default=22)
+@click.option('--identity_file', '-i', required=False, default='')
+@click.option('--password', '-P', required=False, default='')
+@click.option('--forward_ports', '-fp', required=False, multiple=True)
+@click.option('--start_commands', '-sc', required=False, multiple=True)
 @click.option('--config_path', required=False)#, default='')
 def add_server(group_path, alias, host, user, port=22, config_path=CONFIG_PATH, **kwargs):
     """Add a server to a group."""

management/sshman.py

@@ -63,6 +63,13 @@ def add_server(self, group_path, alias, host, user, port, **kwargs):
         """
         last_group = self.make_group(group_path)
 
+        if kwargs.get('forward_ports'):
+            kwargs['forward_ports'] = list(kwargs['forward_ports'])
+
+        if kwargs.get('start_commands'):
+            raise NotImplementedError("Start commands are not supported yet.")
+            kwargs['start_commands'] = list(kwargs['start_commands'])
+
         # Create a server and add it to the last group
         server_data = {'alias': alias,
                        'host': host,
@@ -76,6 +83,4 @@ def add_server(self, group_path, alias, host, user, port, **kwargs):
             server_data['start_commands'] = ['']
 
         server = Server(**server_data)
-        from pprint import pprint
-        pprint(server.dict())
         last_group.add_server(server)

tui/ssh_connections/ssh_connect.py

@@ -50,20 +50,93 @@ def retrieve_file(config_path):
 
     with open(config_path, "r") as json_file:
         config = json.load(json_file)
+        if config.get('user') is None:
+            config['user'] = get_user_from_os()
+
+        if config.get('password'):
+            raise NotImplementedError("Password authentication is not supported yet.")
+
+        if config.get('identity_file'):
+            config['identity_file'] = find_identity_file(config['identity_file'])
+        else:
+            # Declare the default identity file of id_rsa
+            config['identity_file'] = os.path.join(get_user_ssh_dir_from_home(), 'id_rsa')
+
         return config
 
 
+def get_user_from_os():
+    """
+    Attempt to get the user from the OS; assumes that the user you are trying to connect is the same as the user
+    that is invoking the command.
+    :return:
+    """
+    return os.getlogin()
+
+
+def get_user_ssh_dir_from_home():
+    """
+    Gets the SSH directory for the current user.
+    :return:
+    """
+    return os.path.join(os.path.expanduser('~'), '.ssh')
+
+
+def find_identity_file(identity_file):
+    """
+    Finds the identity file from the SSH directory.
+    :param identity_file:
+    :return:
+    """
+    # Determine if the identity file is just the name of a file or if it's a path:
+    if os.path.isabs(identity_file) or identity_file.startswith('~'):
+        if os.path.exists(identity_file):
+            return identity_file
+        else:
+            raise FileNotFoundError(f"Identity file not found: {identity_file}")
+    else:
+        ssh_dir = get_user_ssh_dir_from_home()
+        identity_file_path = os.path.join(ssh_dir, identity_file)
+        if os.path.exists(identity_file_path):
+            return identity_file_path
+        else:
+            raise FileNotFoundError(f"Identity file not found: {identity_file_path}")
+
+def remove_double_spaces(string):
+    """
+    Removes double spaces from a string.
+    :param string:
+    :return:
+    """
+    return ' '.join(string.split())
+
 def connect_shell(config_path):
     config = retrieve_file(config_path)
-    command = f"ssh {config['user']}@{config['host']} -p {config['port']}"
-    print(f'Running: {command}')
+
+    fp = ''
+    if config.get('forward_ports'):
+        for p in config['forward_ports']:
+            fp += f' -L {p} '
+
+    idf = f'-i {config["identity_file"]}'
+
+    command = f"ssh {fp} {idf} {config['user']}@{config['host']} -p {config['port']}"
+    command = remove_double_spaces(command)
+
+    if config.get('start_commands'):
+        raise NotImplementedError("Start commands are not supported yet.")
+        command = [command] + config['start_commands']
+
     return command
 
 
 def connect_sftp(config_path):
     config = retrieve_file(config_path)
-    command = f"sftp -P {config['port']} {config['user']}@{config['host']} "
-    print(f'Running: {command}')
+
+    idf = f"-i {config['identity_file']} "
+
+    command = f"sftp {idf} -P {config['port']} {config['user']}@{config['host']} "
+    command = remove_double_spaces(command)
     return command
 
 

tui/tree.py

@@ -6,9 +6,11 @@
     python code_browser.py PATH
 """
 import json
+import logging
 import subprocess
 from rich.syntax import Syntax
 from rich.traceback import Traceback
+import sys
 
 from textual.app import App, ComposeResult
 from textual.containers import Container, VerticalScroll
@@ -133,4 +135,24 @@ def main():
 
     if getattr(app, 'command', False):
         print(dynamic_data)
-        subprocess.run(dynamic_data, shell=True)
+        # Run the commands in the shell if multiple are specified
+        if isinstance(dynamic_data, list):
+            for i, command in enumerate(dynamic_data):
+                if i == 0:
+                    # Fail if the command does not successfully run
+                    error = subprocess.run(command, shell=True, capture_output=True)
+                    if error.returncode != 0:
+                        ssh_err = f"\nError Connecting to Server:\n{error.stderr.decode('utf-8')}"
+                        print(ssh_err)
+                        sys.exit(error.returncode)
+                else:
+                    error = subprocess.run(command, shell=True)
+                    if error.returncode != 0:
+                        ssh_err = error.stderr.decode('utf-8')
+                        error_msg = f"Error running command: {command} after connecting to the server."
+                        error_msg += '\n{}'.format(ssh_err)
+                        raise RuntimeError(error_msg)
+
+        # Otherwise, just connect:
+        elif isinstance(dynamic_data, str):
+            subprocess.run(dynamic_data, shell=True)