Refactor and simply for GHES/EMU users, add type hints

Author: aegilops

README.md

@@ -26,12 +26,19 @@ You need to be an enterprise administrator to use these scripts!
     pip install -r requirements.txt
     ```
 
-1. Edit the inputs as arguments to the script as follows:
-
-    - the API endpoint (for GHES, EMU, or data residency) in `--api-url`. For GHEC this is not required.
-    - Create a file and save your token there to read it, and call the script with `--token-file` argument, or call the script with the token in `GITHUB_TOKEN` in your environment.
-    - Add the enterprise slug to `--enterprise-slug`. This is string URL version of the enterprise identity.  It's easily available in the enterprise admin url (for cloud and server), e.g. `https://github.com/enterprises/ENTERPRISE-SLUG-HERE`.
-    - For the security manager team script, the list of orgs output by `org-admin-promote.py` in `--unmanaged-orgs` and the name of the security manager team and the team members to add, in `--team-name` and `--team-members`. If you are using GHES 3.15 or below, please use the `--legacy` flag to use the legacy security managers API.
+1. Choose inputs as arguments to the script as follows:
+
+    - the server URL (for GHES, EMU, or data residency) in `--github-url`
+      - For GHEC this is not required.
+    - call the script with the correct GitHub PAT
+      - place it in `GITHUB_TOKEN` in your environment, or
+      - create a file and save your token there to read it, and call the script with the `--token-file` argument
+    - use the enterprise slug as the first argument in the promote/demote scripts
+      - this is string URL version of the enterprise identity.  It's available in the enterprise admin url (for cloud and server), e.g. `https://github.com/enterprises/ENTERPRISE-SLUG-HERE`.
+    - for the security manager team script:
+      - use the list of orgs output by `org-admin-promote.py` in `--unmanaged-orgs`
+      - put the name of the security manager team and the team members to add in `--team-name` and `--team-members`.
+      - If you are using GHES 3.15 or below, please use the `--legacy` flag to use the legacy security managers API.
 
 1. Run them in the following order:
 
@@ -41,31 +48,31 @@ You need to be an enterprise administrator to use these scripts!
 
 ## Assumptions
 
-- The security manager team isn't already an existing team that's using team sync [for enterprise](https://docs.github.com/en/enterprise-cloud@latest/admin/identity-and-access-management/using-saml-for-enterprise-iam/managing-team-synchronization-for-organizations-in-your-enterprise) or [for organizations](https://docs.github.com/en/enterprise-cloud@latest/organizations/organizing-members-into-teams/synchronizing-a-team-with-an-identity-provider-group).  You may be able to edit the script a bit to make this work by adding an existing team to all orgs, but I wasn't going to dive deep into the weeds of identity management.
+- The security manager team isn't already an existing team that's using team sync [for enterprise](https://docs.github.com/en/enterprise-cloud@latest/admin/identity-and-access-management/using-saml-for-enterprise-iam/managing-team-synchronization-for-organizations-in-your-enterprise) or [for organizations](https://docs.github.com/en/enterprise-cloud@latest/organizations/organizing-members-into-teams/synchronizing-a-team-with-an-identity-provider-group).
 
 ## Any extra info?
 
 This is what a successful run looks like.  Here's the inputs:
 
 - The enterprise admin is named `ghe-admin`.
-- The security team is named `spy-stuff` and has two members `luigi` and `hubot`.
+- The security team is named `security-managers` (the default) and has two members `luigi` and `hubot`.
 - The organizations break down as such:
   - `acme` org was already configured correctly.
   - `testorg-00001` needed the team created, with `ghe-admin` removed and `luigi` and `hubot` added.
   - `testorg-00002` was already created
 
 ```console
-$ ./manage-sec-team.py 
-Team spy-stuff updated as a security manager for acme!
-Creating team spy-stuff
-Team spy-stuff updated as a security manager for testorg-00001!
-Removing ghe-admin from spy-stuff
-Adding luigi to spy-stuff
-Adding hubot to spy-stuff
-Creating team spy-stuff
-Team spy-stuff updated as a security manager for testorg-00002!
-Removing ghe-admin from spy-stuff
-Team spy-stuff updated as a security manager for testorg-00003!
+$ ./manage-sec-team.py --sec-team-members luigi hubot
+✓ Team security-managers updated as a security manager for acme
+Creating team security-managers
+✓ Team security-managers updated as a security manager for testorg-00001
+Removing ghe-admin from security-managers
+Adding luigi to security-managers
+Adding hubot to security-managers
+Creating team security-managers
+✓ Team security-managers updated as a security manager for testorg-00002
+Removing ghe-admin from security-managers
+✓ Team security-managers updated as a security manager for testorg-00003
 ```
 
 ## Architecture Footnotes

manage-sec-team.py

@@ -19,29 +19,34 @@
 from argparse import ArgumentParser
 from typing import Any
 from defusedcsv import csv
-from src import teams, organizations, github_token
+from src import teams, organizations, util
 import logging
 
+
 LOG = logging.getLogger(__name__)
 
 
 def add_args(parser) -> None:
     """Add arguments to the command line parser."""
     parser.add_argument(
-        "--api-url",
-        default="https://api.github.com",
-        help="GitHub API URL (https://github-hostname-here/api/v3/ for GHES, EMU or data residency)",
+        "--github-url",
+        required=False,
+        help="GitHub URL for GHES, EMU or data residency",
     )
     parser.add_argument(
         "--token-file",
         required=False,
         help="GitHub Personal Access Token file (or use GITHUB_TOKEN)",
     )
     parser.add_argument(
-        "--org-list", default="all_orgs.csv", help="CSV file of organizations"
+        "--org-list",
+        default="all_orgs.csv",
+        help="CSV file of organizations (default: all_orgs.csv)",
     )
     parser.add_argument(
-        "--sec-team-name", default="security-managers", help="Security team name"
+        "--sec-team-name",
+        default="security-managers",
+        help="Security team name (default: security-managers)",
     )
     parser.add_argument("--sec-team-members", nargs="*", help="Security team members")
     parser.add_argument(
@@ -215,7 +220,7 @@ def main() -> None:
     with open(args.org_list, "r") as f:
         orgs = list(csv.DictReader(f))
 
-    github_pat = github_token.read_token(args.token_file)
+    github_pat = util.read_token(args.token_file)
 
     if not github_pat:
         LOG.error("⨯ GitHub Personal Access Token not found")
@@ -237,6 +242,8 @@ def main() -> None:
         )
         return
 
+    api_url = util.rest_api_url_from_server_url(args.github_url)
+
     # Set up the headers
     headers = {
         "Authorization": "token {}".format(github_pat),
@@ -247,10 +254,10 @@ def main() -> None:
         org_name = org["login"]
 
         make_security_managers_team(
-            org_name, args.sec_team_name, args.api_url, headers, legacy=args.legacy
+            org_name, args.sec_team_name, api_url, headers, legacy=args.legacy
         )
         add_security_managers_to_team(
-            org_name, args.sec_team_name, sec_team_members, args.api_url, headers
+            org_name, args.sec_team_name, sec_team_members, api_url, headers
         )
 
 

org-admin-demote.py

@@ -17,33 +17,34 @@
 
 from argparse import ArgumentParser
 from typing import Iterable, List
-from src import enterprises, github_token
+from src import enterprises, util
 import logging
 
+
 LOG = logging.getLogger(__name__)
 
 
 def add_args(parser: ArgumentParser) -> None:
     """Add arguments to the command line parser."""
     parser.add_argument(
-        "--api-url",
-        default="https://api.github.com/graphql",
-        help="GitHub GraphQL API endpoint (https://github-hostname-here/api/graphql for GHES, EMU or data residency)",
+        "enterprise_slug",
+        help="Enterprise slug (after /enterprises/ in the URL)",
     )
     parser.add_argument(
-        "--token-file",
+        "--github-url",
         required=False,
-        help="File containing a GitHub Personal Access Token with admin:enterprise and read:org scope (or use GITHUB_TOKEN)",
+        help="GitHub URL for GHES, EMU or data residency",
     )
     parser.add_argument(
-        "--enterprise-slug",
-        required=True,
-        help="Enterprise slug (after /enterprises/ in the URL)",
+        "--token-file",
+        required=False,
+        help="File containing a GitHub Personal Access Token with admin:enterprise and read:org scope (or use GITHUB_TOKEN)",
     )
+
     parser.add_argument(
         "--unmanaged-orgs",
         default="unmanaged_orgs.txt",
-        help="Path to newline-delimited list of organization IDs to demote from",
+        help="Path to newline-delimited list of organization IDs to demote from (default: unmanaged_orgs.txt)",
     )
     parser.add_argument(
         "--debug",
@@ -84,7 +85,9 @@ def main() -> None:
 
     logging.basicConfig(level=logging.DEBUG if args.debug else logging.INFO)
 
-    github_pat = github_token.read_token(args.token_file)
+    github_pat = util.read_token(args.token_file)
+
+    api_url = util.graphql_api_url_from_server_url(args.github_url)
 
     if not github_pat:
         LOG.error("⨯ GitHub Personal Access Token not found")
@@ -95,11 +98,11 @@ def main() -> None:
     }
 
     enterprise_id = enterprises.get_enterprise_id(
-        args.api_url, args.enterprise_slug, headers
+        api_url, args.enterprise_slug, headers
     )
 
     unmanaged_orgs = read_unmanaged_org_ids(args.unmanaged_orgs)
-    demote_admin(args.api_url, headers, enterprise_id, unmanaged_orgs)
+    demote_admin(api_url, headers, enterprise_id, unmanaged_orgs)
 
 
 if __name__ == "__main__":  # pragma: no cover

org-admin-promote.py

@@ -18,9 +18,11 @@
 
 from argparse import ArgumentParser
 from typing import List
-from src import enterprises, organizations, github_token
+from urllib.parse import urlparse
+from src import enterprises, organizations, util
 import logging
 
+
 LOG = logging.getLogger(__name__)
 
 
@@ -31,9 +33,9 @@ def add_args(parser: ArgumentParser) -> None:
         help="Enterprise slug (after /enterprises/ in URL)",
     )
     parser.add_argument(
-        "--api-url",
-        default="https://api.github.com/graphql",
-        help="GitHub GraphQL API endpoint (https://github-hostname-here/api/graphql for GHES, EMU or data residency)",
+        "--github-url",
+        required=False,
+        help="GitHub URL for GHES, EMU or data residency",
     )
     parser.add_argument(
         "--token-file",
@@ -43,12 +45,12 @@ def add_args(parser: ArgumentParser) -> None:
     parser.add_argument(
         "--unmanaged-orgs",
         default="unmanaged_orgs.txt",
-        help="Output file for previously unmanaged organization IDs",
+        help="Output file for previously unmanaged organization IDs (default: unmanaged_orgs.txt)",
     )
     parser.add_argument(
         "--orgs-csv",
         default="all_orgs.csv",
-        help="Output CSV file listing all organizations in the enterprise",
+        help="Output CSV file listing all organizations in the enterprise (default: all_orgs.csv)",
     )
     parser.add_argument(
         "--debug",
@@ -119,14 +121,20 @@ def main() -> None:
 
     logging.basicConfig(level=logging.DEBUG if args.debug else logging.INFO)
 
-    github_pat = github_token.read_token(args.token_file)
+    api_url = (
+        util.graphql_api_url_from_server_url(args.github_url)
+        if args.github_url
+        else "https://api.github.com/graphql"
+    )
+
+    github_pat = util.read_token(args.token_file)
     headers: dict[str, str] = {
         "Authorization": f"token {github_pat}",
     }
 
     if (
         promote_all(
-            args.api_url,
+            api_url,
             headers,
             args.enterprise_slug,
             args.unmanaged_orgs,
@@ -137,7 +145,7 @@ def main() -> None:
         return
 
     # Refresh and write all orgs CSV after promotions
-    orgs = organizations.list_orgs(args.api_url, args.enterprise_slug, headers)
+    orgs = organizations.list_orgs(api_url, args.enterprise_slug, headers)
     organizations.write_orgs_to_csv(orgs, args.orgs_csv)
 
 

src/enterprises.py

@@ -4,11 +4,14 @@
 This file holds enterprise-related functions
 """
 
+from typing import Any
 import requests
-from .headers import add_request_headers
+from .util import add_request_headers
 
 
-def get_enterprise_id(api_endpoint, enterprise_slug, headers):
+def get_enterprise_id(
+    api_endpoint: str, enterprise_slug: str, headers: dict[str, str]
+) -> str:
     """
     Get the ID of an enterprise by its slug.
     """
@@ -30,7 +33,7 @@ def get_enterprise_id(api_endpoint, enterprise_slug, headers):
     return response.json()["data"]["enterprise"]["id"]
 
 
-def make_promote_mutation(enterprise_id, org_id, role):
+def make_promote_mutation(enterprise_id: str, org_id: str, role: str) -> str:
     """
     Create a GraphQL mutation to promote the Enterprise owner to owner in an organization.
     """
@@ -49,7 +52,13 @@ def make_promote_mutation(enterprise_id, org_id, role):
     )
 
 
-def promote_admin(api_endpoint, headers, enterprise_id, org_id, role):
+def promote_admin(
+    api_endpoint: str,
+    headers: dict[str, str],
+    enterprise_id: str,
+    org_id: str,
+    role: str,
+) -> dict[str, Any]:
     """
     Promote an enterprise admin to an organization owner.
     """