From 4291d72d7ab91f3442eb406f45d3d8489d721e44 Mon Sep 17 00:00:00 2001
From: James Falcon <therealfalcon@gmail.com>
Date: Tue, 10 Sep 2024 09:11:27 -0500
Subject: [PATCH] feat(ec2): Allow using ec2 profile as creds

---
 VERSION                  |  2 +-
 pycloudlib.toml.template | 13 +++++++++----
 pycloudlib/ec2/cloud.py  | 15 ++++++++++++---
 pycloudlib/ec2/util.py   |  5 ++++-
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/VERSION b/VERSION
index ae66cb8f..a67368f7 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-1!9.0.0
+1!9.1.0
diff --git a/pycloudlib.toml.template b/pycloudlib.toml.template
index 495f8cb4..f18819fa 100644
--- a/pycloudlib.toml.template
+++ b/pycloudlib.toml.template
@@ -24,10 +24,15 @@ tenant_id = ""
 # key_name = ""  # Defaults to your username if not set
 
 [ec2]
-# Most values can be found in ~/.aws/credentials or ~/.aws/config
-access_key_id = ""  # in ~/.aws/credentials
-secret_access_key = ""  # in ~/.aws/credentials
 region = ""  # in ~/.aws/config
+# If 'aws configure sso' has been run, 'profile' should be the only credentials needed
+profile = ""  # in ~/.aws/config
+
+# With modern SSO, these should no longer be necessary.
+# They can be found in ~/.aws/credentials or ~/.aws/config
+# access_key_id = ""  # in ~/.aws/credentials
+# secret_access_key = ""  # in ~/.aws/credentials
+
 # public_key_path = "/root/id_rsa.pub"
 # private_key_path = ""  # Defaults to 'public_key_path' without the '.pub'
 # key_name = ""  # can be found with `aws ec2 describe-key-pairs`
@@ -72,7 +77,7 @@ config_path = "~/.oci/config"
 availability_domain = ""  # Likely in ~/.oci/oci_cli_rc
 compartment_id = ""  # Likely in ~/.oci/oci_cli_rc
 # region = "us-phoenix-1"  # will use region from oci config file if not specified
-# profile = "DEFAULT" # will use default profile from oci config file if not specified 
+# profile = "DEFAULT" # will use default profile from oci config file if not specified
 # public_key_path = "~/.ssh/id_rsa.pub"
 # private_key_path = ""  # Defaults to 'public_key_path' without the '.pub'
 # key_name = ""  # Defaults to your username if not set
diff --git a/pycloudlib/ec2/cloud.py b/pycloudlib/ec2/cloud.py
index 9d1f435b..ea778a13 100644
--- a/pycloudlib/ec2/cloud.py
+++ b/pycloudlib/ec2/cloud.py
@@ -36,6 +36,7 @@ def __init__(
         access_key_id: Optional[str] = None,
         secret_access_key: Optional[str] = None,
         region: Optional[str] = None,
+        profile: Optional[str] = None,
     ):
         """Initialize the connection to EC2.
 
@@ -50,6 +51,7 @@ def __init__(
             access_key_id: user's access key ID
             secret_access_key: user's secret access key
             region: region to login to
+            profile: profile to use from ~/.aws/config
         """
         super().__init__(
             tag,
@@ -59,11 +61,18 @@ def __init__(
         )
         self._log.debug("logging into EC2")
 
+        access_key_id = access_key_id or self.config.get("access_key_id")
+        secret_access_key = secret_access_key or self.config.get(
+            "secret_access_key"
+        )
+        region = region or self.config.get("region")
+        profile = profile or self.config.get("profile")
         try:
             session = _get_session(
-                access_key_id or self.config.get("access_key_id"),
-                secret_access_key or self.config.get("secret_access_key"),
-                region or self.config.get("region"),
+                access_key_id=access_key_id,
+                secret_access_key=secret_access_key,
+                region=region,
+                profile=profile,
             )
             self.client = session.client("ec2")
             self.resource = session.resource("ec2")
diff --git a/pycloudlib/ec2/util.py b/pycloudlib/ec2/util.py
index a380dd49..e82bcd2b 100644
--- a/pycloudlib/ec2/util.py
+++ b/pycloudlib/ec2/util.py
@@ -44,7 +44,9 @@ def _decode_console_output_as_bytes(parsed, **kwargs):
     parsed["OutputBytes"] = base64.b64decode(orig)
 
 
-def _get_session(access_key_id, secret_access_key, region):
+def _get_session(
+    access_key_id=None, secret_access_key=None, region=None, profile=None
+) -> boto3.Session:
     """Get EC2 session.
 
     Args:
@@ -69,4 +71,5 @@ def _get_session(access_key_id, secret_access_key, region):
         aws_access_key_id=access_key_id,
         aws_secret_access_key=secret_access_key,
         region_name=region,
+        profile_name=profile,
     )