GH branch protector

protect.py

import os
import json
from graphqlclient import GraphQLClient

# Set GH_TOKEN environment variable to secret token.
# input on command line.

# Queries and mutations
q_repo = """
query {{
    repository(owner: "{owner}", name: "{name}") {{
        id
    }}
}}
"""
# dirty hack
def q_bpr_pager(owner, name, after=None):
    snip = ""
    if after is not None:
        snip = f", after: \"{after}\""
    return """
query {{
    repository(owner: "{owner}", name: "{name}"{snip}) {{
        branchProtectionRules(first:100) {{
            nodes {{
                pattern
                id
            }}
            pageInfo {{
                endCursor
                hasNextPage
            }}
        }}
    }}
}}
    """.format(owner=owner, name=name, snip=snip)
m_create_bpr = """
mutation {{
    createBranchProtectionRule(input:{{
        repositoryId: "{repoId}",
        pattern: "{pattern}",
        isAdminEnforced: true,
        requiresApprovingReviews: true,
        requiredApprovingReviewCount: 1
    }}) {{
        branchProtectionRule {{
            id
        }}
    }}
}}
"""
m_delete_bpr = """
mutation {{
    deleteBranchProtectionRule(input:{{
        branchProtectionRuleId: "{bprId}"
    }}) {{
        clientMutationId
    }}
}}
"""

def add_bpr(owner, name, pattern):
    repo_res = json.loads(client.execute(q_repo.format(owner=owner, name=name)))
    if "errors" not in repo_res:
        repoId = repo_res["data"]["repository"]["id"]
        res = json.loads(client.execute(m_create_bpr.format(repoId=repoId,
            pattern=pattern)))
        return res
    else:
        return repo_res
def get_bpr(owner, name, pattern):
    page_left = True
    cursor = None
    while page_left:
        res = json.loads(client.execute(q_bpr_pager(owner, name, cursor)))
        if "errors" in res:
            return res
        for node in res["data"]["repository"]["branchProtectionRules"]["nodes"]:
            if node["pattern"] == pattern:
                return json.dumps({"id": node["id"]})
        cursor = res["data"]["repository"]["branchProtectionRules"]["pageInfo"]["endCursor"]
        page_left = res["data"]["repository"]["branchProtectionRules"]["pageInfo"]["hasNextPage"]
    return json.dumps({"errors": ""})

print("Creating client")
client = GraphQLClient("https://api.github.com/graphql")
# Token will be passed as an environment variable
print("Injecting token from GH_TOKEN environment variable")
client.inject_token("Bearer " + os.environ["GH_TOKEN"])
repos = [repo.strip() for repo in input("Repositories to operate on (owner/name), comma separated: ")
    .split(",")]
patterns = [pattern.strip() for pattern in input("Patterns to operate on, comma separated: ")
    .split(",")]
protect = input("Operation (protect/unprotect): ") == "protect"
for repo in repos:
    owner, name = repo.split("/")
    for pattern in patterns:
        if protect:
            print(f"Adding branch protection rule with pattern {pattern} to repository {repo}")
            res = add_bpr(owner, name, pattern)
            if "errors" in res:
                for err in res["errors"]:
                    print(f"Errors occured: {err['message']}")
        else:
            print(f"Removing branch protection rule with pattern {pattern} from repository {repo}")
            bpr_id = json.loads(get_bpr(owner, name, pattern))
            if "errors" in bpr_id:
                print(f"Branch protection rule does not exist, skipping")
            else:
                res = json.loads(client.execute(m_delete_bpr.format(bprId=bpr_id["id"])))
                if "errors" in res:
                    for err in res["errors"]:
                        print(f"Errors occured: {err['message']}")