Merge pull request #11 from arnica-ext/master

Added PR reviews and branch protection

Author: nir-valtman

_config.yml

@@ -1,3 +1,5 @@
+# This is the configuration file for the GitHub Pages.
+# Do not confuse this file with the config.yaml file (without the '_' prefix), which configures GitGoat's behavior.  
 plugins:
   - jekyll-relative-links
 

config.yaml

@@ -23,6 +23,7 @@ repo_names:
 
 ## Specific configurations per repo (optional).
 ### "branch_protection" means the policy is enabled in the main branch
+### "branch_protection_restirctions" define who can push to the protected branch. A team name is specified by the prefix, e.g. the team "Lavender-push" should be included as "push".
 ### "actions_enabled" means that the repo has GitHub Action enabled
 ### "allowed_actions" defines the scope of the allowed actions. The options are "all", "selected" or "local_only". It is applicable only if "actions_enabled" is true.
 ### "verified_allowed_actions" defined if verified actions (not neccesarily owned by GitHub) are allowed. It is applicable only if "actions_enabled" is true.
@@ -33,6 +34,9 @@ repo_configs:
     allowed_actions: all
   Lavender:
     branch_protection: True
+    branch_protection_restirctions:
+      users: [] 
+      teams: ['push']
     actions_enabled: True
     allowed_actions: selected
     verified_allowed_actions: False
@@ -43,6 +47,9 @@ repo_configs:
     verified_allowed_actions: True
   Calendula:
     branch_protection: True
+    branch_protection_restirctions:
+      users: [] 
+      teams: ['push']
     actions_enabled: False
   Tarragon:
     branch_protection: False

run.py

@@ -8,6 +8,7 @@
 from src.members import Membership
 from src.pull_request import PullRequest
 from src.direct_permissions import DirectPermission
+from src.branch import Branch
 
 async def mock(config_file: str, orgs: list = []):
     config = Config() if config_file is None else Config(config_file)
@@ -26,6 +27,8 @@ async def mock(config_file: str, orgs: list = []):
         await create_teams(config, org)
         logging.info('----- Granting Direct Permissions -----')
         await add_direct_permissions(config, org)
+        logging.info('----- Configuring Branch Protection -----')
+        await configure_branch_protection(config, org)
         logging.info('----- Creating Commits and Pull Requests -----')
         await create_commits(config, org)
         logging.info('----- Merging Pull Requests -----')
@@ -89,6 +92,16 @@ async def setup_actions(config, org):
         if config.repo_configs[actions_enabled_repo]['allowed_actions'] == 'selected':
             await a.enable_selected_actions_in_repo(actions_enabled_repo, verified_allowed=config.repo_configs[actions_enabled_repo]['verified_allowed_actions'])
 
+async def configure_branch_protection(config, org):
+    b = Branch(org, config.filename)
+    for repo_name in tqdm(config.repo_names, desc='Branch Protection'):
+        if 'branch_protection_restirctions' in config.repo_configs[repo_name]:
+            users = config.repo_configs[repo_name]['branch_protection_restirctions']['users']
+            teams = []
+            for team_postfix in config.repo_configs[repo_name]['branch_protection_restirctions']['teams']:
+                teams.append(f'{repo_name}-{team_postfix}')
+            await b.set_branch_protection(repo_name, 'main', users, teams)
+
 async def create_commits(config, org):
     r = Repository(org, config.filename)
     pr = PullRequest(org, config.filename)
@@ -97,23 +110,41 @@ async def create_commits(config, org):
             token =  member['token'] if 'ghp_' in member['token'] else 'ghp_' + member['token']
             repo = await r.clone(commit_details['repo'], member['login'], token, member['email'], commit_details['branch'])
             c = Commit(repo)
-            c.generate_commits(200, commit_details['days'])
+            c.generate_commits(50, commit_details['days'])
             if commit_details['create_pr']:
                 await pr.create_pull_request(token, commit_details['repo'], commit_details['branch'])
                 #logging.info(f'Created a PR by {member["login"]} from branch {commit_details["branch"]}')
 
 async def merge_pull_requests(config, org):
     pr = PullRequest(org, config.filename)
+    reviewed_prs = {}
     for member in tqdm(config.members, desc=f'Member PRs'):
         for repo in config.repo_names:
-            if f'{repo}-triage' in member['member_of_groups'] or f'{repo}-maintain' in member['member_of_groups']:
+            reviewed_prs[repo] = []
+            if is_member_allowed_to_merge(config, member, repo):
                 token =  member['token'] if 'ghp_' in member['token'] else 'ghp_' + member['token']
                 prs = await pr.get_pull_requests(token,repo)
                 for p in prs:
-                    merged = await pr.merge(token, repo, p)
-                    if not merged: 
-                        logging.warning(f'Did NOT merge the PR id {p} in repository {repo} by {member["login"]}')
+                    if prs[p] != member['login']:
+                        await pr.review(token, repo, p)
+                        reviewed_prs[repo].append(p)
+                    if p in reviewed_prs[repo]:
+                        merged = await pr.merge(token, repo, p)
+                        if not merged: 
+                            logging.warning(f'Did NOT merge the PR id {p} in repository {repo} by {member["login"]}')
 
+def is_member_allowed_to_merge(config, member, repo):
+    if 'branch_protection_restirctions' in config.repo_configs[repo]:
+        if member['login'] in config.repo_configs[repo]['branch_protection_restirctions']['users']:
+            return True
+        for team in config.repo_configs[repo]['branch_protection_restirctions']['teams']:
+            for group_membership in member['member_of_groups']:
+                if group_membership == f'{repo}-{team}':
+                    return True
+    else:
+        if f'{repo}-maintain' in member['member_of_groups'] or f'{repo}-push' in member['member_of_groups']:
+            return True
+    return False
 
 if __name__ == '__main__':
     try:

src/branch.py

@@ -4,6 +4,7 @@ class Branch:
 
     def __init__(self, organization, config_file = None):
         self.endpoint = f'/repos/{organization}/[REPO]/git/refs'
+        self.branch_protection_endpoint = f'/repos/{organization}/[REPO]/branches/[BRANCH]/protection'
         self.config_file = config_file
 
     async def get_main(self, pat, repository):
@@ -21,4 +22,31 @@ async def create_branch(self, pat, repository, branch_name, source_branch_sha):
         }
         resp = await conn.post(endpoint, json_data=data)
         return resp
+    
+    async def set_branch_protection(self, repository, branch_name, restricted_users = [], restricted_teams = []):
+        conn = ConnectionHandler(config_file=self.config_file)
+        endpoint = self.branch_protection_endpoint.replace('[REPO]',repository).replace('[BRANCH]', branch_name)
+        payload = {
+            'required_status_checks': {
+                'strict': False,
+                'contexts': []
+            },
+            'enforce_admins': False,
+            'require_code_owner_reviews': False,
+            'required_pull_request_reviews': {
+                'required_approving_review_count': 1
+            },
+            'restrictions': {
+                'users': restricted_users,
+                'teams': restricted_teams
+            },
+            'allow_force_pushes': True,
+            'allow_deletions': True,
+            'bypass_pull_request_allowances': {
+                'users': restricted_users,
+                'teams': restricted_teams
+            }
+        }
+        resp = await conn.put(endpoint, payload)
+        return resp
 

src/pull_request.py

@@ -11,11 +11,11 @@ def __init__(self, organization, config_file = None):
     async def get_pull_requests(self, pat, repository):
         conn = ConnectionHandler(pat, self.config_file)
         endpoint = self.endpoint.replace('[REPO]', repository)
-        pr_ids = []
+        pr_ids = {}
         resp = await conn.get(endpoint)
         for pr in resp:
             if pr['state'] == 'open': 
-                pr_ids.append(pr['number'])
+                pr_ids[pr['number']] = pr['user']['login']
         return pr_ids
 
     async def create_pull_request(self, pat, repository, head_branch):
@@ -30,6 +30,16 @@ async def create_pull_request(self, pat, repository, head_branch):
         resp = await conn.post(endpoint, json_data=data)
         return resp['number']
 
+    async def review(self, pat, repository, pull_request_number):
+        conn = ConnectionHandler(pat, self.config_file)
+        endpoint = self.endpoint.replace('[REPO]', repository) + f'/{str(pull_request_number)}/reviews'
+        data = {
+            'event': 'APPROVE',
+            'body:': self.fake.lexify(text='GitGoat automated PR review ????????')
+        }
+        resp = await conn.post(endpoint, data)
+        return resp
+
     async def merge(self, pat, repository, pull_request_number):
         conn = ConnectionHandler(pat, self.config_file)
         endpoint = self.endpoint.replace('[REPO]', repository) + f'/{str(pull_request_number)}/merge'