works.py

#!/usr/bin/env python3
"""
generate_readme_v2.py — Professional AWS Accounts Directory Generator (v2)

Generates a rich, visually polished README.md with:
  - HTML tables with color-coded rows (prod/sandbox/dev)
  - Clickable switch-role buttons styled as badges
  - Summary statistics dashboard
  - Collapsible OU tree
  - Versatile (external org) accounts section
  - Table of contents & quick-jump anchors
"""

import boto3
import time
import threading
from datetime import datetime, timezone
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  CONFIGURATION
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

# Map SSO usernames → display names. Add/remove as needed.
SSO_MAPPING = {
    # "john.doe": "John Doe",
    # "jane.smith": "Jane Smith",
}

# IAM switch-role names
SWITCH_ROLES = ["ca-iam-cie-engineer", "ca-iam-cie-a"]

# AWS Switch Role console URL template
SWITCH_ROLE_URL = (
    "https://signin.aws.amazon.com/switchrole"
    "?roleName={role}&account={account_id}&displayName={display}"
)

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  VERSATILE ACCOUNTS (not part of our Org — maintained manually)
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
VERSATILE_ACCOUNTS = [
    {
        "metadata_name": "versatile-test",
        "account_id": "021891588050",
        "account_name": "test",
        "versatile_ou": "Root/Workloads",
    },
    # {
    #     "metadata_name": "versatile-network",
    #     "account_id": "XXXXXXXXXXXX",
    #     "account_name": "network",
    #     "versatile_ou": "Root/Workloads",
    # },
]

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  TAG KEYS
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
TAG_DESCRIPTION = "syf:aws:account_description"
TAG_PRIMARY_OWNER = "syf:aws:account.primary_owner"
TAG_SECONDARY_OWNER = "syf:aws:account.secondary_owner"
TAG_ALLOWED_CIS = "syf:aws:account.allowed_cis"
TAG_AVM_RESOURCE = "syf:aws:avm_resource"

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  THREADING / RATE-LIMIT CONFIGURATION
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

# AWS Organizations API rate limits:
#   - ListTagsForResource : 1 TPS  (but bursty up to ~5)
#   - ListParents         : 1 TPS
#   - DescribeOU          : 1 TPS
# We use a conservative concurrency of 5 workers with a token-bucket
# throttle that caps the *overall* Organizations API call rate to
# ~8 requests/sec across all threads.  This stays well within burst
# limits while being ~5x faster than serial execution.

MAX_WORKERS = 5                 # parallel threads
API_CALLS_PER_SECOND = 8.0     # global rate cap (all threads combined)

# ── Global token-bucket rate limiter ─────────────────────────────────

class _TokenBucket:
    """Simple token-bucket rate limiter (thread-safe)."""

    def __init__(self, rate: float, burst: int = 1):
        self._rate = rate
        self._burst = burst
        self._tokens = float(burst)
        self._last = time.monotonic()
        self._lock = threading.Lock()

    def acquire(self):
        while True:
            with self._lock:
                now = time.monotonic()
                self._tokens = min(
                    self._burst,
                    self._tokens + (now - self._last) * self._rate,
                )
                self._last = now
                if self._tokens >= 1.0:
                    self._tokens -= 1.0
                    return
            time.sleep(1.0 / self._rate)

_rate_limiter = _TokenBucket(rate=API_CALLS_PER_SECOND, burst=5)

# ── Thread-local boto3 clients (boto3 clients are NOT thread-safe) ──

_thread_local = threading.local()

def _get_org_client():
    """Return a per-thread Organizations client."""
    if not hasattr(_thread_local, "org_client"):
        _thread_local.org_client = boto3.client("organizations")
    return _thread_local.org_client

# ── OU name cache (many accounts share the same parent OUs) ─────────

_ou_cache = {}          # ou_id → ou_name
_ou_cache_lock = threading.Lock()

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  COLOR / STYLE CONFIGURATION
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

# Row background colors (light tints for readability)
ENV_STYLES = {
    "prod":    {"bg": "#fde8e8", "badge_bg": "#dc2626", "badge_text": "#fff", "label": "PROD",    "emoji": "🔴"},
    "sandbox": {"bg": "#e6f9e6", "badge_bg": "#16a34a", "badge_text": "#fff", "label": "SANDBOX", "emoji": "🟢"},
    "other":   {"bg": "#e8f0fe", "badge_bg": "#2563eb", "badge_text": "#fff", "label": "DEV/LAB", "emoji": "🔵"},
}


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  HELPERS
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def get_all_active_accounts(org_client):
    accounts = []
    paginator = org_client.get_paginator("list_accounts")
    for page in paginator.paginate():
        for acct in page["Accounts"]:
            if acct["Status"] == "ACTIVE":
                accounts.append(acct)
    return sorted(accounts, key=lambda a: a["Name"].lower())


def get_account_tags(org_client, account_id):
    _rate_limiter.acquire()
    tags = {}
    paginator = org_client.get_paginator("list_tags_for_resource")
    for page in paginator.paginate(ResourceId=account_id):
        for tag in page["Tags"]:
            tags[tag["Key"]] = tag["Value"]
    return tags


def _resolve_ou_name(org_client, ou_id):
    """Return OU name, using a shared cache to avoid duplicate API calls."""
    with _ou_cache_lock:
        if ou_id in _ou_cache:
            return _ou_cache[ou_id]

    _rate_limiter.acquire()
    ou = org_client.describe_organizational_unit(
        OrganizationalUnitId=ou_id
    )["OrganizationalUnit"]
    name = ou["Name"]

    with _ou_cache_lock:
        _ou_cache[ou_id] = name
    return name


def get_parent_ous(org_client, account_id):
    parts = []
    child_id = account_id
    while True:
        _rate_limiter.acquire()
        resp = org_client.list_parents(ChildId=child_id)
        parent = resp["Parents"][0]
        parent_id = parent["Id"]
        parent_type = parent["Type"]
        if parent_type == "ROOT":
            parts.append("Root")
            break
        else:
            parts.append(_resolve_ou_name(org_client, parent_id))
            child_id = parent_id
    parts.reverse()
    ou_path = "/".join(parts)
    ou_name = parts[-1] if len(parts) > 1 else "Root"
    return ou_name, ou_path


def format_owner(sso_value):
    if not sso_value or sso_value.strip() == "":
        return "<em>N/A</em>"
    sso = sso_value.strip()
    name = SSO_MAPPING.get(sso)
    return f"{sso} ({name})" if name else sso


def classify_env(ou_path):
    """Return 'prod', 'sandbox', or 'other' based on OU path."""
    path_lower = ou_path.lower()
    if "prod" in path_lower:
        return "prod"
    if "sandbox" in path_lower:
        return "sandbox"
    return "other"


def make_switch_role_badge(account_id, account_name, role, color, text_color="#fff"):
    """Return an HTML anchor styled as a colored badge/button."""
    display = f"{account_name}+{role}"
    url = SWITCH_ROLE_URL.format(role=role, account_id=account_id, display=display)
    short_label = role.replace("ca-iam-cie-", "")  # shorter label for badge
    return (
        f'<a href="{url}" title="Switch to {role}">'
        f'<img src="https://img.shields.io/badge/{short_label}-{role}-{color.strip("#")}?style=flat-square&labelColor=333" alt="{role}"/>'
        f"</a>"
    )


def make_switch_role_cell(account_id, account_name):
    """Two shield badges per account."""
    badges = []
    palette = ["2563eb", "7c3aed"]  # blue, purple
    for i, role in enumerate(SWITCH_ROLES):
        badges.append(
            make_switch_role_badge(account_id, account_name, role, palette[i % len(palette)])
        )
    return " ".join(badges)


def make_env_badge_html(env_key):
    """Inline HTML badge for environment."""
    style = ENV_STYLES[env_key]
    return (
        f'<span style="background:{style["badge_bg"]};color:{style["badge_text"]};'
        f'padding:2px 8px;border-radius:12px;font-size:11px;font-weight:700;'
        f'letter-spacing:.5px;">{style["label"]}</span>'
    )


def build_ou_tree_text(accounts_with_ou):
    """Build a text-based OU tree."""
    tree = {}
    for item in accounts_with_ou:
        parts = item["ou_path"].split("/")
        node = tree
        for p in parts:
            node = node.setdefault(p, {})

    lines = []

    def _render(node, prefix=""):
        keys = sorted(node.keys())
        for i, key in enumerate(keys):
            is_last = i == len(keys) - 1
            connector = "└── " if is_last else "├── "
            lines.append(f"{prefix}{connector}{key}")
            extension = "    " if is_last else "│   "
            _render(node[key], prefix + extension)

    root_keys = sorted(tree.keys())
    for rk in root_keys:
        lines.append(rk)
        _render(tree[rk])

    return "\n".join(lines)


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  MARKDOWN / HTML BUILDERS
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def render_header(now, total, env_counts, avm_count, legacy_count):
    return f"""\
<div align="center">

# ☁️ AWS Accounts Directory

<p>
  <img src="https://img.shields.io/badge/Accounts-{total}-0078D4?style=for-the-badge&logo=amazonaws&logoColor=white" alt="Total Accounts"/>
  <img src="https://img.shields.io/badge/Production-{env_counts.get('prod',0)}-dc2626?style=for-the-badge" alt="Prod"/>
  <img src="https://img.shields.io/badge/Sandbox-{env_counts.get('sandbox',0)}-16a34a?style=for-the-badge" alt="Sandbox"/>
  <img src="https://img.shields.io/badge/Dev%2FLab-{env_counts.get('other',0)}-2563eb?style=for-the-badge" alt="Dev/Lab"/>
  <img src="https://img.shields.io/badge/AVM-{avm_count}-6d28d9?style=for-the-badge" alt="AVM"/>
  <img src="https://img.shields.io/badge/Legacy-{legacy_count}-a16207?style=for-the-badge" alt="Legacy"/>
</p>

<p><em>Centralized inventory of all active AWS accounts in our organization.</em></p>
<p><sub>🕒 Last generated: <strong>{now}</strong></sub></p>

</div>

---
"""


def render_toc():
    return """\
## 📑 Table of Contents

| Section | Description |
|---------|-------------|
| [Active Accounts](#-active-accounts) | All accounts in our AWS Organization |
| [Versatile Accounts](#-versatile-accounts-external-org) | External org accounts with switch-role access |
| [OU Tree](#-organization-ou-tree) | Organizational Unit hierarchy |
| [Legend](#-legend) | Color coding & column descriptions |

---
"""


def render_legend():
    return """\
## 📖 Legend

### Environment Colors

| Color | Environment | Row Background | Meaning |
|:-----:|-------------|----------------|---------|
| 🔴 | **Production** | <span style="background:#fde8e8;padding:2px 12px;">Light Red</span> | Production workloads — handle with care |
| 🟢 | **Sandbox** | <span style="background:#e6f9e6;padding:2px 12px;">Light Green</span> | Sandbox / experimentation accounts |
| 🔵 | **Dev / Lab** | <span style="background:#e8f0fe;padding:2px 12px;">Light Blue</span> | Development, lab, and other environments |

### Column Reference

| Column | Description |
|--------|-------------|
| **Account Name** | Friendly name of the AWS account |
| **Account ID** | 12-digit AWS account identifier |
| **Switch Role** | Click badges to assume role in target account via AWS Console |
| **OU Name** | Immediate parent Organizational Unit |
| **OU Path** | Full path from Root to the account's OU |
| **AVM / Legacy** | `AVM` = Account Vending Machine provisioned · `Legacy` = manually created |
| **Allowed CIs** | Application CIs authorized for this account |
| **Primary Owner** | SSO of the primary account owner |
| **Secondary Owner** | SSO of the secondary account owner |
| **Description** | Purpose / description from account tags |

---
"""


def render_accounts_table(rows):
    """Build an HTML table with colored rows."""
    header = """\
## ☁️ Active Accounts

<table>
<thead>
<tr>
  <th align="center">#</th>
  <th align="center">Env</th>
  <th>Account Name</th>
  <th>Account ID</th>
  <th align="center">Switch Role</th>
  <th>OU Name</th>
  <th>OU Path</th>
  <th align="center">AVM / Legacy</th>
  <th align="center">Allowed CIs</th>
  <th>Primary Owner</th>
  <th>Secondary Owner</th>
  <th>Description</th>
</tr>
</thead>
<tbody>
"""
    body_lines = []
    for idx, r in enumerate(rows, start=1):
        style = ENV_STYLES[r["env"]]
        bg = style["bg"]
        env_badge = make_env_badge_html(r["env"])
        avm_badge = (
            '<code style="background:#6d28d9;color:#fff;padding:2px 6px;border-radius:4px;">AVM</code>'
            if r["avm_legacy"] == "AVM"
            else '<code style="background:#a16207;color:#fff;padding:2px 6px;border-radius:4px;">Legacy</code>'
        )

        body_lines.append(f"""\
<tr style="background:{bg};">
  <td align="center"><strong>{idx}</strong></td>
  <td align="center">{env_badge}</td>
  <td><strong>{r['name']}</strong></td>
  <td><code>{r['id']}</code></td>
  <td align="center">{r['switch_role']}</td>
  <td>{r['ou_name']}</td>
  <td><code>{r['ou_path']}</code></td>
  <td align="center">{avm_badge}</td>
  <td align="center">{r['allowed_cis']}</td>
  <td>{r['primary_owner']}</td>
  <td>{r['secondary_owner']}</td>
  <td>{r['description']}</td>
</tr>""")

    footer = """\
</tbody>
</table>

---
"""
    return header + "\n".join(body_lines) + "\n" + footer


def render_versatile_table():
    if not VERSATILE_ACCOUNTS:
        return ""

    header = """\
## 🔗 Versatile Accounts (External Org)

> **Note:** These accounts are **not** part of our AWS Organization.
> We have cross-account switch-role access into them.

<table>
<thead>
<tr>
  <th align="center">#</th>
  <th>Metadata Name</th>
  <th>Account ID</th>
  <th>Account Name</th>
  <th align="center">Switch Role</th>
  <th>Versatile OU</th>
</tr>
</thead>
<tbody>
"""
    body_lines = []
    for idx, v in enumerate(VERSATILE_ACCOUNTS, start=1):
        sr = make_switch_role_cell(v["account_id"], v["account_name"])
        body_lines.append(f"""\
<tr>
  <td align="center"><strong>{idx}</strong></td>
  <td><strong>{v['metadata_name']}</strong></td>
  <td><code>{v['account_id']}</code></td>
  <td>{v['account_name']}</td>
  <td align="center">{sr}</td>
  <td><code>{v['versatile_ou']}</code></td>
</tr>""")

    footer = """\
</tbody>
</table>

---
"""
    return header + "\n".join(body_lines) + "\n" + footer


def render_ou_tree(tree_text):
    return f"""\
## 🌳 Organization OU Tree

<details>
<summary><strong>Click to expand full OU hierarchy</strong></summary>

```
{tree_text}
```

</details>

---
"""


def render_footer(now):
    return f"""\
<div align="center">
<sub>
  Generated by <code>generate_readme_v2.py</code> &nbsp;|&nbsp; {now} &nbsp;|&nbsp; ☁️ Cloud Infrastructure Engineering
</sub>
</div>
"""


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  MAIN
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def _process_account(acct, progress_counter, total):
    """
    Worker function executed in a thread pool.
    Each thread gets its own boto3 client via _get_org_client().
    All API calls go through the shared _rate_limiter.
    """
    org_client = _get_org_client()
    aid = acct["Id"]
    name = acct["Name"]

    tags = get_account_tags(org_client, aid)
    ou_name, ou_path = get_parent_ous(org_client, aid)

    avm_tag = tags.get(TAG_AVM_RESOURCE, "").strip().lower()
    is_avm = avm_tag not in ("", "false")
    env = classify_env(ou_path)

    with progress_counter["lock"]:
        progress_counter["done"] += 1
        done = progress_counter["done"]
    print(f"   [{done}/{total}] {name} ({aid})")

    return {
        "name": name,
        "id": aid,
        "switch_role": make_switch_role_cell(aid, name),
        "ou_name": ou_name,
        "ou_path": ou_path,
        "env": env,
        "avm_legacy": "AVM" if is_avm else "Legacy",
        "allowed_cis": tags.get(TAG_ALLOWED_CIS, "N/A").strip() or "N/A",
        "primary_owner": format_owner(tags.get(TAG_PRIMARY_OWNER, "")),
        "secondary_owner": format_owner(tags.get(TAG_SECONDARY_OWNER, "")),
        "description": tags.get(TAG_DESCRIPTION, "N/A").strip() or "N/A",
    }


def main():
    org_client = boto3.client("organizations")

    print("⏳ Fetching active accounts …")
    accounts = get_all_active_accounts(org_client)
    total = len(accounts)
    print(f"   Found {total} active accounts.")
    print(f"   Using {MAX_WORKERS} threads with {API_CALLS_PER_SECOND} API calls/sec cap.\n")

    progress = {"done": 0, "lock": threading.Lock()}
    rows_map = {}   # account_id → row dict (to preserve original sort order)

    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
        futures = {
            pool.submit(_process_account, acct, progress, total): acct["Id"]
            for acct in accounts
        }
        for future in as_completed(futures):
            aid = futures[future]
            try:
                rows_map[aid] = future.result()
            except Exception as exc:
                print(f"   ⚠️  Failed to process {aid}: {exc}")

    # Restore original alphabetical order
    rows = [rows_map[acct["Id"]] for acct in accounts if acct["Id"] in rows_map]

    # ── Statistics ───────────────────────────────────────────────────
    env_counts = Counter(r["env"] for r in rows)
    avm_count = sum(1 for r in rows if r["avm_legacy"] == "AVM")
    legacy_count = len(rows) - avm_count
    ou_tree_text = build_ou_tree_text(rows)
    now = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")

    # ── Assemble README ─────────────────────────────────────────────
    readme = ""
    readme += render_header(now, len(rows), env_counts, avm_count, legacy_count)
    readme += render_toc()
    readme += render_legend()
    readme += render_accounts_table(rows)
    readme += render_versatile_table()
    readme += render_ou_tree(ou_tree_text)
    readme += render_footer(now)

    with open("README.md", "w") as f:
        f.write(readme)

    print(f"\n✅ README.md generated — {len(rows)} accounts, {len(VERSATILE_ACCOUNTS)} versatile.")
    print(f"   OU cache hits saved ~{len(_ou_cache)} redundant DescribeOU calls.")


if __name__ == "__main__":
    main()