#!/usr/bin/env python
# coding: utf-8

# # Bulk-add branch protection rules
# 
# Going through all jupyterhub repos and adding the most basic branch protection rule for `main`

# I use a token stored in my keychain.
# I have two of these, one read-only and one with full access.
# 
# You can create similar tokens with:
# 
# - visit https://github.com/settings/personal-access-tokens/new
# - store with:
# 
#   ```
#   python -m keyring set api.github.com read-only
#   ```

# In[1]:


import requests
import keyring
s = requests.Session()
s.headers['Authorization'] = f"bearer {keyring.get_password('api.github.com', 'full-access')}"
s.headers['X-Github-Next-Global-ID'] = '1'


# In[3]:


s.get("https://api.github.com/user").json()['login']


# ## Step 1: collect repos
# 
# Use graphQL endpoint to collect repositories
# 
# 1. only non-forks
# 2. only need:
#    - default branch name
#    - if there is an existing branch protection rule
#    - repo's global id for use in mutation later

# In[12]:


from jinja2 import Template
import json
from tqdm.notebook import tqdm

query_template = Template("""
{
  organization(login:"jupyterhub") {
#  user(login: "minrk") {
    repositories(isFork: false, ownerAffiliations: OWNER, first: 100{%- if after %}, after: "{{ after }}"{%- endif %}) {
      totalCount
      pageInfo {
        endCursor
        hasNextPage
      }
      nodes {
        id
        name
        defaultBranchRef {
          name
        }
        branchProtectionRules(first: 5) {
          nodes {
            pattern
            matchingRefs(first: 5) {
              nodes {
                name
              }
            }
          }
        }
      }
    }
  }
}
""")

github_graphql = "https://api.github.com/graphql"

def fetch_repos():
    
    after = None
    repos = []
    has_next_page = True
    progress = tqdm(desc="fetching", unit="repos")
    while has_next_page:
        r = s.post(github_graphql, data=json.dumps(dict(query=query_template.render(after=after))))
        r.raise_for_status()
        response = r.json()["data"]["organization"]["repositories"]
        progress.total = response["totalCount"]
        progress.update(len(response['nodes']))
        repos.extend(response['nodes'])
        after = response['pageInfo']['endCursor']
        has_next_page = response['pageInfo']['hasNextPage']
    progress.close()
    return repos
repos = fetch_repos()


# In[6]:


repos[0]


# ## Step 2: select which repos to protect
# 
# Iterate through repos, finding:
# 
# 1. if default branch is not main, note that it still needs a rename
# 2. if the default branch is not protected, add repo to list of repos to protect

# In[13]:


to_protect = []
no_main = []
is_protected = []

for repo in repos:
    print(repo['name'])
    if repo['defaultBranchRef'] is None:
        print("  Empty repo!")
        continue
    default_branch = repo['defaultBranchRef']['name']
    if default_branch != 'main':
        print(f"  non-main default branch: {default_branch}")
        no_main.append(repo)

    rules = repo['branchProtectionRules']['nodes']
    matched = False
    for rule in rules:
        if default_branch in [ref['name'] for ref in rule['matchingRefs']['nodes']]:
            print(f" Has rule: {rule['pattern']}")
            matched = True
            is_protected.append(repo)
            break

    if not matched:
        print(f"  Needs rule")
        to_protect.append(repo)

print(f"{len(repos)} total repos")
print(f"{len(no_main)} repos missing main branch")
print(f"{len(is_protected)} repos with protected main")
print(f"{len(to_protect)} repos will get new protection rules")


# In[14]:


mutation_template = Template("""
mutation {
  {% for repo in repos %}
  {{ repo['name'] | replace("-", "") | replace(".","") }}: createBranchProtectionRule(input: {repositoryId: "{{ repo['id'] }}", pattern: "{{ repo['defaultBranchRef']['name'] }}"}) {
    clientMutationId
  }
  {% endfor %}
}
""") 
print(mutation_template.render(repos=to_protect)[:1024])


# In[15]:


r = s.post(github_graphql, data=json.dumps(dict(query=mutation_template.render(repos=to_protect))))
r.raise_for_status()


# In[16]:


r.json()