add retarget-sso-session

local-app/python-tools/aws_config_editor/aws_config_editor.py

@@ -12,6 +12,12 @@
   edit    – Sed-style find/replace on key *values* inside matching profiles.
   create  – Create a new profile from inline key=value pairs, an existing
             profile clone, or a Jinja2 template file.
+  retarget-sso-session
+          – Repair profiles generated by aws-sso-util (which emits a bare
+            sso_start_url per profile) to instead reference a named
+            [sso-session] stanza via sso_session.  Removes sso_start_url,
+            sso_account_id-adjacent redundant URL fields, and inserts
+            sso_session = <name> in each affected profile.
 
 Comment & whitespace preservation
 ----------------------------------
@@ -58,6 +64,10 @@
       --template profiles/govcloud_role.j2 \\
       --var "account_id=777788889999" --var "role=AdminRole"
 
+  # Retarget aws-sso-util-generated profiles to use sso_session (dry-run first)
+  python aws_config_editor.py retarget-sso-session --dry-run
+  python aws_config_editor.py retarget-sso-session --yes
+
   # Operate on credentials file instead of config
   python aws_config_editor.py --file ~/.aws/credentials list
 """
@@ -75,7 +85,7 @@
 
 # ── Version ───────────────────────────────────────────────────────────────────
 
-__version__ = "1.0.1"
+__version__ = "1.0.2"
 
 # ── Defaults ──────────────────────────────────────────────────────────────────
 
@@ -138,6 +148,15 @@ def sub_key(self, key: str, search: str, replace: str, regex: bool) -> bool:
                 return False
         return False
 
+    def delete_key(self, key: str) -> bool:
+        """Remove all lines for *key* from this profile. Returns True if any were removed."""
+        before = len(self.lines)
+        self.lines = [
+            rl for rl in self.lines
+            if not re.match(rf"^\s*{re.escape(key)}\s*=", rl.text)
+        ]
+        return len(self.lines) < before
+
     def clone(self, new_name: str, is_config: bool) -> "Profile":
         """Return a deep copy with a new section header."""
         new_header = _make_section_header(new_name, is_config)
@@ -197,6 +216,22 @@ def get(self, name: str) -> Profile | None:
     def matching(self, pattern: str, substring: bool) -> list[Profile]:
         return [p for p in self.profiles if _matches(p.name, pattern, substring)]
 
+    def sso_sessions(self) -> dict[str, str]:
+        """
+        Return a mapping of {sso_session_name: sso_start_url} for every
+        [sso-session NAME] stanza found in the file.
+        """
+        result: dict[str, str] = {}
+        for prof in self.profiles:
+            # Section header is the raw text, e.g. '[sso-session my-sso]'
+            m = re.match(r"^\[sso-session\s+(.+)\]\s*$", prof.section_header.strip())
+            if m:
+                session_name = m.group(1).strip()
+                url = prof.keys().get("sso_start_url", "")
+                if url:
+                    result[session_name] = url
+        return result
+
     # ── mutation ──────────────────────────────────────────────────────────────
 
     def remove(self, name: str) -> bool:
@@ -464,6 +499,169 @@ def cmd_create(args: argparse.Namespace) -> int:
     return 0
 
 
+def cmd_retarget_sso_session(args: argparse.Namespace) -> int:
+    """
+    Repair profiles generated by aws-sso-util to reference a named
+    [sso-session] stanza instead of carrying a bare sso_start_url.
+
+    aws-sso-util populate-profiles writes profiles like:
+
+        [profile dev-account-ReadOnly]
+        sso_start_url  = https://my-sso.awsapps.com/start
+        sso_account_id = 111122223333
+        sso_role_name  = ReadOnly
+        sso_region     = us-east-1
+        region         = us-east-1
+        output         = json
+
+    The AWS CLI v2 prefers profiles that reference a shared [sso-session]
+    stanza instead:
+
+        [sso-session my-sso]
+        sso_start_url  = https://my-sso.awsapps.com/start
+        sso_region     = us-east-1
+        sso_registration_scopes = sso:account:access
+
+        [profile dev-account-ReadOnly]
+        sso_session    = my-sso
+        sso_account_id = 111122223333
+        sso_role_name  = ReadOnly
+        region         = us-east-1
+        output         = json
+
+    This command:
+      1. Reads every [sso-session NAME] stanza and records its sso_start_url.
+      2. Finds every [profile X] whose sso_start_url matches one of those URLs.
+      3. In each matching profile:
+           - Removes  sso_start_url  (now owned by the session stanza)
+           - Removes  sso_region     (now owned by the session stanza)
+           - Inserts  sso_session = NAME  as the first key after the header
+
+    Profiles that already have sso_session, or that have an sso_start_url not
+    covered by any [sso-session] stanza, are left untouched and reported.
+
+    Options
+    -------
+      --pattern / --substring  Restrict which profiles are eligible (optional).
+                               Profiles outside the pattern are always skipped.
+      --keep-sso-region        Do NOT remove sso_region from retargeted profiles
+                               (use if your profiles need a per-profile override).
+    """
+    cf = _open_file(args)
+
+    # ── Step 1: collect sso-session stanzas ───────────────────────────────────
+    sessions = cf.sso_sessions()   # {session_name: sso_start_url}
+
+    if not sessions:
+        print("[ERROR] No [sso-session] stanzas found in the config file.", file=sys.stderr)
+        print("        Create one first, e.g.:", file=sys.stderr)
+        print("          [sso-session my-sso]", file=sys.stderr)
+        print("          sso_start_url  = https://my-sso.awsapps.com/start", file=sys.stderr)
+        print("          sso_region     = us-east-1", file=sys.stderr)
+        print("          sso_registration_scopes = sso:account:access", file=sys.stderr)
+        return 1
+
+    print(f"Found {len(sessions)} [sso-session] stanza(s):")
+    for name, url in sessions.items():
+        print(f"  [{name}]  →  {url}")
+    print()
+
+    # Build reverse map: url → session_name  (last one wins if duplicates exist)
+    url_to_session: dict[str, str] = {url: name for name, url in sessions.items()}
+
+    # ── Step 2: identify candidate profiles ───────────────────────────────────
+    # Start with all profiles, then optionally restrict by --pattern.
+    candidates = (
+        cf.matching(args.pattern, args.substring)
+        if args.pattern
+        else cf.profiles
+    )
+
+    # Categorise each candidate
+    # (profile, session_name) for profiles that will be retargeted
+    to_retarget: list[tuple[Profile, str]] = []
+    already_done: list[str] = []
+    url_not_matched: list[tuple[str, str]] = []   # (profile_name, url)
+    no_sso_url: list[str] = []
+
+    for prof in candidates:
+        # Skip [sso-session] stanzas themselves
+        if re.match(r"^\[sso-session", prof.section_header.strip()):
+            continue
+
+        kv = prof.keys()
+
+        if "sso_session" in kv:
+            already_done.append(prof.name)
+            continue
+
+        if "sso_start_url" not in kv:
+            # Not an SSO profile at all — silently skip
+            continue
+
+        url = kv["sso_start_url"]
+        if url in url_to_session:
+            to_retarget.append((prof, url_to_session[url]))
+        else:
+            url_not_matched.append((prof.name, url))
+
+    # ── Step 3: report ────────────────────────────────────────────────────────
+    if already_done:
+        print(f"Skipped ({len(already_done)} already have sso_session):")
+        for n in already_done:
+            print(f"  {n}")
+        print()
+
+    if url_not_matched:
+        print(f"Skipped ({len(url_not_matched)} sso_start_url not covered by any sso-session):")
+        for n, url in url_not_matched:
+            print(f"  {n}  →  {url}")
+        print()
+
+    if not to_retarget:
+        print("Nothing to retarget.")
+        return 0
+
+    print(f"{len(to_retarget)} profile(s) will be retargeted:")
+    for prof, session_name in to_retarget:
+        kv = prof.keys()
+        keys_removed = ["sso_start_url"]
+        if not args.keep_sso_region and "sso_region" in kv:
+            keys_removed.append("sso_region")
+        print(f"  {prof.name}")
+        print(f"    + sso_session    = {session_name}")
+        for k in keys_removed:
+            print(f"    - {k}  (was: {kv[k]})")
+
+    if args.dry_run:
+        print("\nDRY RUN — no changes written.")
+        return 0
+
+    if not _confirm(f"\nRetarget {len(to_retarget)} profile(s)?", args.yes):
+        print("Aborted.")
+        return 0
+
+    print(f"Backup: {cf.backup()}")
+
+    # ── Step 4: apply mutations ───────────────────────────────────────────────
+    for prof, session_name in to_retarget:
+        kv = prof.keys()
+
+        # Remove sso_start_url (and optionally sso_region) first
+        prof.delete_key("sso_start_url")
+        if not args.keep_sso_region:
+            prof.delete_key("sso_region")
+
+        # Insert sso_session as the first key (line index 1, after the header)
+        # Only insert if it doesn't already exist (guard against re-runs)
+        if "sso_session" not in prof.keys():
+            prof.lines.insert(1, RawLine(f"sso_session = {session_name}"))
+
+    cf.write()
+    print(f"\nRetargeted {len(to_retarget)} profile(s).")
+    return 0
+
+
 # ── Jinja2 helpers (optional) ─────────────────────────────────────────────────
 
 def _render_jinja(text: str, variables: dict[str, str]) -> str:
@@ -602,6 +800,22 @@ def build_parser() -> argparse.ArgumentParser:
     _add_mutation_args(cp)
     cp.set_defaults(func=cmd_create)
 
+    # retarget-sso-session
+    rp = sub.add_parser(
+        "retarget-sso-session",
+        help="Replace sso_start_url in profiles with sso_session references (aws-sso-util repair)",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description=cmd_retarget_sso_session.__doc__,
+    )
+    _add_pattern_args(rp)
+    rp.add_argument(
+        "--keep-sso-region",
+        action="store_true",
+        help="Do not remove sso_region from retargeted profiles",
+    )
+    _add_mutation_args(rp)
+    rp.set_defaults(func=cmd_retarget_sso_session)
+
     return parser