ssm-cli 0.1.4.dev0__tar.gz → 0.1.4.dev3__tar.gz

{ssm_cli-0.1.4.dev0/ssm_cli.egg-info → ssm_cli-0.1.4.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ssm-cli
-Version: 0.1.4.dev0
+Version: 0.1.4.dev3
 Summary: CLI tool to help with SSM functionality, aimed at adminstrators
 Author-email: Simon Fletcher <simon.fletcher@lexisnexisrisk.com>
 License: MIT License

{ssm_cli-0.1.4.dev0 → ssm_cli-0.1.4.dev3}/ssm_cli/aws.py

@@ -19,10 +19,10 @@ _session_cache = []
 _client_cache = {}
 
 @contextlib.contextmanager
-def aws_session():
+def aws_session(use_cache=True):
     """ A context manager for creating a boto3 session with caching built in """
     try:
-        if len(_session_cache) > 0:
+        if len(_session_cache) > 0 and use_cache:
             yield _session_cache[0]
             return
         
@@ -42,10 +42,10 @@ def aws_session():
         raise e
 
 @contextlib.contextmanager
-def aws_client(service_name):
+def aws_client(service_name, use_cache=True):
     """ A context manager for creating a boto3 client with caching built in """
-    with aws_session() as session:
-        if service_name in _client_cache:
+    with aws_session(use_cache) as session:
+        if service_name in _client_cache and use_cache:
             yield _client_cache[service_name]
             return
         

{ssm_cli-0.1.4.dev0 → ssm_cli-0.1.4.dev3}/ssm_cli/cli.py

@@ -1,37 +1,34 @@
 import sys
+import logging
 
 from rich_argparse import ArgumentDefaultsRichHelpFormatter
 
 import confclasses
 from ssm_cli.config import CONFIG
-from ssm_cli.xdg import get_log_file, get_conf_file
+from ssm_cli.xdg import get_conf_file, get_log_file
 from ssm_cli.commands import COMMANDS
 from ssm_cli.cli_args import CliArgumentParser, ARGS
 from ssm_cli.aws import AWSAuthError, AWSAccessDeniedError
 from ssm_cli.console import console
+from ssm_cli.logging import setup_logging, configure_log_level
+from ssm_cli.instances import cleanup_sessions
 from rich.markup import escape
 
-# Setup logging
-import logging
-logging.basicConfig(
-    level=logging.WARNING,
-    filename=get_log_file(),
-    filemode='+wt',
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-    datefmt='%Y-%m-%d %H:%M:%S'
-)
 logger = logging.getLogger(__name__)
 
 def cli(argv: list = None) -> int:
     if argv is None:
         argv = sys.argv[1:]
 
+    # Initialize logging
+    setup_logging()
+
     # Manually set the log level now, so we get accurate logging during argument parsing
     for i, arg in enumerate(argv):
         if arg == '--log-level':
-            logging.getLogger().setLevel(argv[i+1].upper())
+            configure_log_level(argv[i+1])
         if arg.startswith('--log-level='):
-            logging.getLogger().setLevel(arg.split('=')[1].upper())
+            configure_log_level(arg.split('=')[1])
 
     logger.debug(f"CLI called with {argv}")
 
@@ -68,12 +65,12 @@ def cli(argv: list = None) -> int:
         console.print(f"Invalid config: {e}", style="red")
         return 1
     
-    logging.getLogger().setLevel(CONFIG.log.level.upper())
+    configure_log_level(CONFIG.log.level)
 
     
     for logger_name, level in CONFIG.log.loggers.items():
         logger.debug(f"setting logger {logger_name} to {level}")
-        logging.getLogger(logger_name).setLevel(level.upper())
+        configure_log_level(level, name=logger_name)
 
 
     if ARGS.command not in COMMANDS:
@@ -103,3 +100,5 @@ def _run():
         console.print(f"Error: {e}", style="red bold")
         logger.exception(e, stack_info=True, stacklevel=20)
         return 1
+    finally:
+        cleanup_sessions()

{ssm_cli-0.1.4.dev0 → ssm_cli-0.1.4.dev3}/ssm_cli/commands/setup.py

@@ -1,7 +1,7 @@
 import argparse
 import paramiko
 from ssm_cli.commands.base import BaseCommand
-from ssm_cli.xdg import get_conf_root, get_conf_file, get_ssh_hostkey, get_log_file
+from ssm_cli.xdg import get_conf_root, get_conf_file, get_ssh_hostkey
 from ssm_cli.cli_args import ARGS
 from ssm_cli.config import CONFIG
 from confclasses import from_dict, save

{ssm_cli-0.1.4.dev0 → ssm_cli-0.1.4.dev3}/ssm_cli/commands/ssh_proxy/__init__.py

@@ -2,7 +2,7 @@ import socket
 import time
 
 from ssm_cli.commands.ssh_proxy.server import SshServer
-from ssm_cli.instances import Instance, Instances
+from ssm_cli.instances import Instance, Instances, SessionManagerPluginError, SessionManagerPluginPortError
 from ssm_cli.config import CONFIG
 from ssm_cli.commands.base import BaseCommand
 from ssm_cli.cli_args import ARGS
@@ -50,9 +50,14 @@ def direct_tcpip_callback(instance: Instance, connections: dict) -> callable:
                 try:
                     internal_port = get_free_port()
                     proc = instance.start_port_forwarding_session_to_remote_host(host, remote_port, internal_port)
-                except Exception as e: # TODO: catch more specific exceptions
-                    logger.warning(f"session-manager-plugin attempt {attempt} failed: {e}")
+                    connections[(host, remote_port)] = (internal_port, proc)
+                    break
+                except SessionManagerPluginPortError as e:
+                    logger.warning(f"session-manager-plugin attempt {attempt} failed due to port clash, retrying: {e}")
                     time.sleep(0.1)
+                except SessionManagerPluginError as e:
+                    logger.error(f"session-manager-plugin failed: {e}")
+                    return None
         
         logger.debug(f"connecting to session manager plugin on 127.0.0.1:{internal_port}")
         # Even though we wait for the process to say its connected, we STILL need to wait for it
@@ -63,6 +68,7 @@ def direct_tcpip_callback(instance: Instance, connections: dict) -> callable:
                     raise RuntimeError("session-manager-plugin has exited")
                 sock = socket.create_connection(('127.0.0.1', internal_port))
                 logger.info(f"connected to 127.0.0.1:{internal_port}")
+                break
             except Exception as e:
                 logger.warning(f"connection attempt {attempt} failed: {e}")
                 time.sleep(0.1)

{ssm_cli-0.1.4.dev0 → ssm_cli-0.1.4.dev3}/ssm_cli/commands/ssh_proxy/server.py

@@ -81,8 +81,11 @@ class SshServer(paramiko.ServerInterface):
         logger.info(f"direct TCP/IP request: chan={chanid} origin={origin} destination={destination}")
         host = destination[0]
         remote_port = destination[1]
-
-        sock = self.direct_tcpip_callback(host, remote_port)
+        try:
+            sock = self.direct_tcpip_callback(host, remote_port)
+        except Exception as e:
+            logger.error(f"failed to connect to session manager plugin: {e}")
+            return paramiko.OPEN_FAILED_CONNECT_FAILED
         
         if not sock:
             logger.error("failed to connect to session manager plugin")

{ssm_cli-0.1.4.dev0 → ssm_cli-0.1.4.dev3}/ssm_cli/instances.py

@@ -1,4 +1,3 @@
-from dataclasses import dataclass, field
 from functools import cache
 import json
 import re
@@ -14,6 +13,38 @@ from ssm_cli.aws import aws_client, aws_session
 import logging
 logger = logging.getLogger(__name__)
 
+class SessionManagerPluginError(Exception):
+    """ A generic exception for any AWS errors """
+    stdout = []
+    returncode = 0
+    def __init__(self, message, stdout, returncode):
+        super().__init__(message)
+        self.stdout = stdout
+        self.returncode = returncode
+    
+    def __str__(self):
+        return f"{super().__str__()} (returncode={self.returncode}, stdout={self.stdout})"
+
+class SessionManagerPluginPortError(SessionManagerPluginError):
+    """ A specific exception for a timeout error """
+    pass
+
+_sessions = {}
+
+def cleanup_sessions():
+    """ Kill any active sessions """
+    with aws_client('ssm') as client:
+        for session_id, proc in _sessions.items():
+            logger.debug(f"terminating session {session_id}")
+            try:
+                proc.terminate()
+                proc.wait()
+                client.terminate_session(SessionId=session_id)
+            except Exception as e:
+                logger.error(f"Failed to cleanup session {session_id} but continuing anyway: {e}")
+
+    _sessions.clear()
+
 
 class Instance:
     """
@@ -56,7 +87,8 @@ class Instance:
         
     def start_port_forwarding_session_to_remote_host(self, host: str, remote_port: int, internal_port: int):
         logger.debug(f"start port forwarding between localhost:{internal_port} and {host}:{remote_port} via {self.id}")
-        with aws_session() as session, aws_client('ssm') as client:
+        with aws_session(False) as session:
+            client = session.client('ssm') # we need a fresh connection where session is the same as the client
 
             parameters = dict(
                 Target=self.id,
@@ -88,15 +120,22 @@ class Instance:
                 stdout=subprocess.PIPE,
                 stderr=subprocess.STDOUT
             )
-
-            for line in proc.stdout: # this will block, needs sorting
+            _sessions[response["SessionId"]] = proc
+
+            # changes needed here:
+            # SessionManagerPluginPortError needs to be raised correctly
+            # add a timeout for waiting on stdout
+            # maybe split up stdout/stderr?
+            output = b''
+            for line in proc.stdout:
                 line = line.strip()
+                output += line + b'\n'
                 if line == b'Waiting for connections...':
                     return proc
                 else:
                     logger.debug(f"Recieved from session-manager-plugin: {line}")
-
-            raise RuntimeError("Failed to start port forwarding session")
+            
+            raise SessionManagerPluginError("Failed to start port forwarding session", output, proc.returncode)
 
 def _session_manager_plugin( command: list) -> int:
     """ Call out to subprocess and ignore interrupts """

ssm_cli-0.1.4.dev3/ssm_cli/logging.py

@@ -0,0 +1,41 @@
+import logging
+import threading
+from datetime import datetime, timedelta
+from typing import Optional
+from ssm_cli.xdg import get_log_file, get_all_log_files
+
+
+def setup_logging():
+    """Set up basic logging configuration with date-based rotation."""
+    logging.basicConfig(
+        level=logging.WARNING,
+        filename=get_log_file(),
+        filemode='a',
+        format='%(asctime)s - %(process)d [%(threadName)s] - %(name)s - %(levelname)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S'
+    )
+    start_log_cleanup()
+
+
+def configure_log_level(level: str, name: Optional[str] = None):
+    """Configure logger level."""
+    logging.getLogger(name).setLevel(level.upper())
+
+
+def cleanup_old_logs(days_to_keep: int = 7):
+    """Clean up log files older than specified days. Runs with no error handling."""
+    cutoff_date = datetime.now() - timedelta(days=days_to_keep)
+    
+    for log_file in get_all_log_files():
+        # Extract date from filename
+        date_part = log_file.stem.split('.')[-1]
+        file_date = datetime.strptime(date_part, '%Y-%m-%d')
+        
+        if file_date < cutoff_date:
+            log_file.unlink()
+
+
+def start_log_cleanup(days_to_keep: int = 7):
+    """Start log cleanup in background thread without error handling."""
+    thread = threading.Thread(target=cleanup_old_logs, args=(days_to_keep,), daemon=True)
+    thread.start()

{ssm_cli-0.1.4.dev0 → ssm_cli-0.1.4.dev3}/ssm_cli/xdg.py

@@ -1,5 +1,7 @@
 from pathlib import Path
 from xdg_base_dirs import xdg_config_home, xdg_state_home
+from datetime import datetime
+from typing import Optional
 
 # Default locations for ssm-cli files
 # ~/.config/ssm-cli/ssm.yaml
@@ -24,10 +26,16 @@ def get_conf_file(check=True) -> Path:
         raise EnvironmentError(f"{path} missing, run `ssm setup` to create")
     return path
 
-def get_log_file() -> Path:
-    path = get_state_root() / 'run.log'
+def get_log_file(date: Optional[datetime] = None) -> Path:
+    if date is None:
+        date = datetime.now()
+    date_str = date.strftime('%Y-%m-%d')
+    path = get_state_root() / f"run.{date_str}.log"
     return path
 
+def get_all_log_files() -> str:
+    return get_state_root().glob('run.*.log')
+
 def get_ssh_hostkey(check=True) -> Path:
     path = get_conf_root(check) / 'hostkey.pem'
     if check and not path.exists():

{ssm_cli-0.1.4.dev0 → ssm_cli-0.1.4.dev3/ssm_cli.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ssm-cli
-Version: 0.1.4.dev0
+Version: 0.1.4.dev3
 Summary: CLI tool to help with SSM functionality, aimed at adminstrators
 Author-email: Simon Fletcher <simon.fletcher@lexisnexisrisk.com>
 License: MIT License

{ssm_cli-0.1.4.dev0 → ssm_cli-0.1.4.dev3}/ssm_cli.egg-info/SOURCES.txt

@@ -9,6 +9,7 @@ ssm_cli/cli_args.py
 ssm_cli/config.py
 ssm_cli/console.py
 ssm_cli/instances.py
+ssm_cli/logging.py
 ssm_cli/xdg.py
 ssm_cli.egg-info/PKG-INFO
 ssm_cli.egg-info/SOURCES.txt