refactor: sdk and cli commands

Author: containerscrew

src/cloudsnake/cli/cli.py

@@ -39,7 +39,6 @@ def version_cmd():
         )
     )
 
-
 @app.callback()
 def entrypoint(
     ctx: typer.Context,

src/cloudsnake/cli/ssm.py

@@ -8,12 +8,7 @@
 
 EC2_RUNNING_FILTER = "Name=instance-state-name,Values=running"
 
-EC2_INSTANCE_SELECTOR_QUERY = (
-    "Reservations[*].Instances[*].{"
-    "TargetId: InstanceId, "
-    "Name: Tags[?Key=='Name'].Value | [0]"
-    "}"
-)
+EC2_INSTANCE_SELECTOR_QUERY = "[].{TargetId: InstanceId, Name: Tags[?Key=='Name'].Value | [0]}"
 
 ssm = typer.Typer(
     no_args_is_help=True,
@@ -39,18 +34,20 @@ def start_session(
     ),
 ):
     ssm = SSMStartSessionWrapper(
+        session=ctx.obj.session,
         profile=ctx.obj.profile,
         region=ctx.obj.region,
     )
-    ssm.create_client(ctx.obj.session)
+
     if with_instance_selector:
         ec2 = EC2InstanceWrapper(
+            session=ctx.obj.session,
             filters=EC2_RUNNING_FILTER,
             query=EC2_INSTANCE_SELECTOR_QUERY,
             profile=ctx.obj.profile,
             region=ctx.obj.region,
         )
-        ec2.create_client(ctx.obj.session)
+
         instances = ec2.describe_ec2_instances()
         if not instances:
             typer.echo("No running instances found.")

src/cloudsnake/sdk/aws.py

@@ -1,48 +1,51 @@
+from __future__ import annotations
+
 import logging
-from botocore.config import Config
+from abc import ABC, abstractmethod
+from typing import Optional, Any
 
+import boto3
+from botocore.config import Config
 
-class App:
-    """
-    Common AWS wrapper base class for cloudsnake SDK.
-    Provides: logging, client creation, session storage, region/profile propagation.
-    """
 
+class App(ABC):
     def __init__(
         self,
-        client=None,
-        filters=None,
-        query=None,
-        profile=None,
-        region=None,
-        session=None,
+        session: Optional[boto3.Session] = None,
+        region: Optional[str] = None,
+        profile: Optional[str] = None,
+        filters: Optional[dict] = None,
+        query: Optional[Any] = None,
+        retries: int = 10,
         **kwargs,
     ):
-        self.log = logging.getLogger("cloudsnake")
-
+        self.log = logging.getLogger(self.__class__.__name__)
         self.filters = filters
         self.query = query
-        self.client_name = client
+        self.session = session
         self.profile = profile
         self.region = region
-        self.session = session
-        self.client = None
+        self.retries = retries
+        self._client = None
 
-    def create_client(self, session):
-        """
-        Create a boto3 client using the provided session.
-        Stores the session, region, and profile internally.
-        """
-        self.session = session
-        try:
-            if not self.region:
-                self.region = session.region_name
-            if not self.profile and hasattr(session, "profile_name"):
-                self.profile = session.profile_name
-        except Exception:
-            pass
-
-        config = Config(retries={"max_attempts": 10, "mode": "standard"})
-
-        self.client = session.client(self.client_name, config=config)
-        return self.client
+    @property
+    @abstractmethod
+    def client_name(self) -> str:
+        pass
+
+    @property
+    def client(self):
+        if self._client is None:
+            if not self.session:
+                raise RuntimeError("No boto3 session available")
+            self._client = self._create_client()
+        return self._client
+
+    def _create_client(self):
+        if not self.region and self.session.region_name:
+            self.region = self.session.region_name
+        if not self.profile and getattr(self.session, "profile_name", None):
+            self.profile = self.session.profile_name
+
+        cfg = Config(retries={"max_attempts": self.retries, "mode": "standard"})
+        return self.session.client(self.client_name, config=cfg)

src/cloudsnake/sdk/boto3_session.py

@@ -1,59 +1,41 @@
+from __future__ import annotations
+
+import boto3
 import logging
 import os
-import boto3
+from typing import Optional
 
 
 class SessionWrapper:
-    """Encapsulates AWS boto3 Session operations"""
-
-    def __init__(self, profile: str | None = None, region: str = "us-east-1"):
-        """
-        :param profile: AWS credentials profile (defaults to value in AWS_PROFILE).
-        :param region: AWS region.
-        """
+    def __init__(self, profile: Optional[str] = None, region: str = "us-east-1"):
         self.log = logging.getLogger("cloudsnake.session")
-
         self.profile = profile or os.getenv("AWS_PROFILE")
         self.region = region
 
         if not self.profile:
-            self.log.warning("No AWS profile provided. Using environment default.")
+            self.log.warning("No AWS profile provided, falling back to environment defaults")
 
-        self.log.debug(
-            f"SessionWrapper initialized with profile={self.profile} region={self.region}"
-        )
+        self.log.debug(f"SessionWrapper(profile={self.profile}, region={self.region})")
 
     def with_local_session(self) -> boto3.Session:
-        """Return a boto3 Session using ~/.aws/credentials"""
-        self.log.debug("Using local AWS credentials via ~/.aws/credentials")
-        return boto3.Session(
-            profile_name=self.profile,
-            region_name=self.region,
-        )
+        self.log.debug("Creating local boto3.Session using ~/.aws/credentials")
+        return boto3.Session(profile_name=self.profile, region_name=self.region)
 
     def with_sts_assume_role_session(self, role_arn: str) -> boto3.Session:
-        """Return a boto3 Session via STS AssumeRole"""
         self.log.debug(f"Assuming role via STS: {role_arn}")
 
-        base_session = boto3.Session(
-            profile_name=self.profile,
-            region_name=self.region,
-        )
-        sts = base_session.client("sts")
-
-        response = sts.assume_role(
-            RoleArn=role_arn,
-            RoleSessionName="cloudsnake-session",
-        )
+        base = boto3.Session(profile_name=self.profile, region_name=self.region)
+        sts = base.client("sts")
 
-        creds = response["Credentials"]
+        res = sts.assume_role(RoleArn=role_arn, RoleSessionName="cloudsnake-session")
+        creds = res["Credentials"]
 
-        assumed = boto3.Session(
+        session = boto3.Session(
             aws_access_key_id=creds["AccessKeyId"],
             aws_secret_access_key=creds["SecretAccessKey"],
             aws_session_token=creds["SessionToken"],
             region_name=self.region,
         )
 
-        self.log.debug("AssumeRole session created successfully")
-        return assumed
+        self.log.debug("STS assumed role session created")
+        return session

src/cloudsnake/sdk/ec2.py

@@ -1,54 +1,43 @@
+from __future__ import annotations
+
 import jmespath
+from typing import Any, Dict, List, Optional
+
 from cloudsnake.helpers import parse_filters
 from cloudsnake.sdk.aws import App
 from botocore.exceptions import ClientError
 
 
 class EC2InstanceWrapper(App):
-    """
-    Wrapper class for managing EC2 instances.
-    """
-
-    def __init__(self, client="ec2", filters=None, query=None, **kwargs):
-        """
-        Initialize the EC2 class.
-
-        Args:
-            client (EC2Client): The EC2 client object.
-            instances (list, optional): A list of EC2 instances. Defaults to None.
-            filters (str, optional): Additional filters for querying instances. Defaults to "".
-            query (str, optional): A query string for filtering instances. Defaults to "".
-            **kwargs: Additional keyword arguments.
-
-        """
-        super().__init__(client, filters, query, **kwargs)
-        self.instances = {}
-
-    def describe_ec2_instances(self):
-        """
-        AWS EC2 describe instances.
-        """
-        if self.filters:
-            parsed_filters = parse_filters(self.filters)
-        else:
-            parsed_filters = []
-
-        self.log.info("Describing EC2 instances")
+    def __init__(
+        self,
+        filters: Optional[str] = None,
+        query: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(filters=filters, query=query, **kwargs)
+        self.instances: Dict[str, Any] = {}
+
+    @property
+    def client_name(self) -> str:
+        return "ec2"
+
+    def describe_ec2_instances(self) -> Any:
+        parsed_filters = parse_filters(self.filters) if self.filters else []
+
         try:
             paginator = self.client.get_paginator("describe_instances")
-
             for page in paginator.paginate(Filters=parsed_filters):
-                self.instances.update(page)
+                for reservation in page.get("Reservations", []):
+                    for instance in reservation.get("Instances", []):
+                        iid = instance.get("InstanceId")
+                        if iid:
+                            self.instances[iid] = instance
+
+            return jmespath.search(self.query, list(self.instances.values())) if self.query else list(self.instances.values())
 
-            if self.query is not None:
-                result = jmespath.search(self.query, self.instances)
-                return result
-            else:
-                return self.instances
         except ClientError as err:
             self.log.error(
-                "Couldn't register device",
-                err.response["Error"]["Code"],
-                err.response["Error"]["Message"],
+                f"EC2 describe_instances failed: {err.response['Error']['Code']} - {err.response['Error']['Message']}"
             )
             raise

src/cloudsnake/sdk/ssm_session.py

@@ -1,53 +1,49 @@
+from __future__ import annotations
+
 import errno
 import json
 import shutil
 import subprocess
 import logging
+from typing import Optional, Dict, Any
 
 from cloudsnake.helpers import ignore_user_entered_signals
 from cloudsnake.sdk.aws import App
 
 
 PLUGIN_NOT_FOUND_MSG = """
 Session Manager Plugin not found.
-
-Install it following:
 https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html
 """
 
 
 class SSMStartSessionWrapper(App):
-    """Encapsulates Amazon SSM Start Session actions."""
-
-    def __init__(self, client="ssm", session_response_output=None, **kwargs):
-        super().__init__(client, **kwargs)
+    def __init__(
+        self,
+        session_response_output: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
         self.session_response_output = session_response_output
         self.log = logging.getLogger("cloudsnake.ssm")
 
+    @property
+    def client_name(self) -> str:
+        return "ssm"
+
     def _ensure_plugin_installed(self):
         if shutil.which("session-manager-plugin") is None:
-            self.log.error("SessionManagerPlugin not found")
             raise FileNotFoundError(PLUGIN_NOT_FOUND_MSG)
 
-    def start_session_response(self, target: str) -> dict:
-        """Start an SSM session and store the response."""
-        self.log.debug(f"Calling ssm.start_session(Target='{target}')")
-
-        response = self.client.start_session(
-            Target=target,
-            Reason="Session started by cloudsnake",
-        )
-        self.session_response_output = response
-        return response
+    def start_session_response(self, target: str) -> Dict[str, Any]:
+        self.log.debug(f"ssm.start_session(Target={target})")
+        res = self.client.start_session(Target=target, Reason="cloudsnake session")
+        self.session_response_output = res
+        return res
 
     def start_session(self, target: str):
-        """
-        Start an SSM session using the session-manager-plugin.
-        Uses: self.profile and self.region inherited from App.
-        """
         self._ensure_plugin_installed()
-
-        self.log.info(f"Starting SSM session for instance {target}...")
+        self.log.info(f"Starting SSM session for {target}")
         self.start_session_response(target)
 
         try:
@@ -63,28 +59,26 @@ def start_session(self, target: str):
                         f"https://ssm.{self.region}.amazonaws.com",
                     ]
                 )
-            self.log.info("SSM session closed normally")
+            self.log.info("Session closed cleanly")
             return 0
 
         except subprocess.CalledProcessError as e:
-            self.log.error(f"Failed to start SSM session: {e}", exc_info=True)
+            self.log.error(f"SSM session failed: {e}")
             self.terminate_session()
             raise
 
         except OSError as ex:
             if ex.errno == errno.ENOENT:
-                self.log.error("SessionManagerPlugin is missing", exc_info=True)
                 self.terminate_session()
                 raise FileNotFoundError(PLUGIN_NOT_FOUND_MSG) from ex
             else:
                 self.log.error("OS error during SSM session", exc_info=True)
                 raise
 
     def terminate_session(self) -> None:
-        """Terminate the SSM session."""
         if self.session_response_output and "SessionId" in self.session_response_output:
-            session_id = self.session_response_output["SessionId"]
-            self.log.debug(f"Terminating SSM session {session_id}")
-            self.client.terminate_session(SessionId=session_id)
+            sid = self.session_response_output["SessionId"]
+            self.log.debug(f"terminate_session {sid}")
+            self.client.terminate_session(SessionId=sid)
         else:
-            self.log.warning("No active session to terminate")
+            self.log.warning("No session to terminate")