Merge pull request #28 from jitsecurity/bigquery

Bigquery

Author: JonathanRosenboim

.gitignore

@@ -19,3 +19,4 @@ venv*/*
 .ipynb_checkpoints
 .DS_*
 node_modules/
+teams*

requirements.txt

@@ -1,5 +1,7 @@
 pydantic==2.1.1
 PyGithub==1.59.1
 python-dotenv==1.0.0
+pytz==2024.1
+google-cloud-bigquery==3.19.0
 requests==2.31.0
 loguru==0.7.0

src/scripts/sync_teams/sync_teams.py

@@ -21,23 +21,31 @@
 logger.add(sys.stderr, format=logger_format)
 
 
-def parse_input_file() -> Organization:
+def parse_input_file() -> Tuple[Organization, bool, bool]:
     """
-    Parse the input JSON file and return an Organization object.
+    Parse the input JSON file and return an Organization object, the skip_no_resources flag,
+    and the verify_github_membership flag.
 
     Returns:
-        Organization: The parsed organization object.
+        Tuple[Organization, bool, bool]: The parsed organization object, the skip_no_resources flag,
+        and the verify_github_membership flag.
     """
-    # Create the argument parser
     parser = argparse.ArgumentParser(description="Retrieve teams and assets")
-
-    # Add the file argument
     parser.add_argument("file", help="Path to a JSON file")
 
-    # Parse the command line arguments
+    # Default behavior is True, override with --no-skip-no-resources and --no-verify-github-membership
+    parser.add_argument("--skip-no-resources", dest='skip_no_resources', action="store_true",
+                        help="Skip teams with no active resources", default=True)
+    parser.add_argument("--no-skip-no-resources", dest='skip_no_resources', action="store_false",
+                        help="Do not skip teams with no active resources")
+
+    parser.add_argument("--verify-github-membership", dest='verify_github_membership', action="store_true",
+                        help="Verify GitHub membership when setting team members", default=True)
+    parser.add_argument("--no-verify-github-membership", dest='verify_github_membership', action="store_false",
+                        help="Do not verify GitHub membership when setting team members")
+
     args = parser.parse_args()
 
-    # Check if the file exists and is a JSON file
     if not os.path.isfile(args.file):
         logger.error("Error: File does not exist.")
         sys.exit(1)
@@ -46,14 +54,13 @@ def parse_input_file() -> Organization:
         sys.exit(1)
     logger.info(f"Reading file: {args.file}")
 
-    # Read the JSON file
     with open(args.file, "r") as file:
         json_data = file.read()
-
-    # Parse the JSON data
+    logger.info(f"JSON data: {json_data}")
     try:
         data = json.loads(json_data)
-        return Organization(teams=[TeamStructure(**team) for team in data["teams"]])
+        return Organization(teams=[TeamStructure(**team) for team in data["teams"]]), \
+            args.skip_no_resources, args.verify_github_membership
     except (ValidationError, KeyError) as e:
         logger.error(f"Failed to validate input file: {e}")
         sys.exit(1)
@@ -125,14 +132,15 @@ def get_teams_to_delete(topic_names: List[str], existing_team_names: List[str])
     return get_different_items_in_lists(existing_team_names, topic_names)
 
 
-def get_desired_teams(assets: List[Asset], organization: Organization) -> List[str]:
+def get_desired_teams(assets: List[Asset], organization: Organization, skip_no_resources: bool) -> List[str]:
     """
     Get the desired teams based on the assets and organization.
     Also filter out teams that match the TEAM_WILDCARD_TO_EXCLUDE environment variable.
 
     Args:
         assets (List[Asset]): The list of assets.
         organization (Organization): The organization object.
+        skip_no_resources (bool): Whether to skip teams with no active resources.
 
     Returns:
         List[str]: The names of the desired teams.
@@ -143,7 +151,7 @@ def get_desired_teams(assets: List[Asset], organization: Organization) -> List[s
         for resource in team.resources:
             if resource.type == ResourceType.GithubRepo and resource.name in [asset.asset_name for asset in assets]:
                 team_resources.append(resource.name)
-        if team_resources:
+        if team_resources or not skip_no_resources:
             desired_teams.append(team.name)
         else:
             logger.info(
@@ -165,7 +173,8 @@ def get_desired_teams(assets: List[Asset], organization: Organization) -> List[s
 
 
 def process_teams(token, organization, assets: List[Asset],
-                  existing_teams: List[TeamAttributes]) -> Tuple[List[str], List[TeamAttributes]]:
+                  existing_teams: List[TeamAttributes],
+                  skip_no_resources: bool) -> Tuple[List[str], List[TeamAttributes]]:
     """
     Process the teams in the organization and create or delete teams as necessary.
     We will delete the teams at a later stage to avoid possible synchronization issues.
@@ -174,13 +183,13 @@ def process_teams(token, organization, assets: List[Asset],
         token (str): The JWT token.
         organization (Organization): The organization object.
         existing_teams (List[TeamAttributes]): The existing teams.
-
+        skip_no_resources (bool): Whether to skip teams with no active resources.
     Returns:
         Tuple[List[str], List[TeamAttributes]]: The names of the teams to delete and the created teams.
     """
     logger.info("Determining required changes in teams.")
 
-    desired_teams = get_desired_teams(assets, organization)
+    desired_teams = get_desired_teams(assets, organization, skip_no_resources)
     existing_team_names = [team.name for team in existing_teams]
     teams_to_create = get_teams_to_create(desired_teams, existing_team_names)
     teams_to_delete = get_teams_to_delete(desired_teams, existing_team_names)
@@ -193,7 +202,7 @@ def process_teams(token, organization, assets: List[Asset],
 
 
 def process_members(token: str, organization: Organization, existing_teams: List[TeamAttributes],
-                    desired_teams: List[str]) -> None:
+                    desired_teams: List[str], verify_github_membership: bool) -> None:
     logger.info("Processing team members.")
     for team_structure in organization.teams:
         try:
@@ -210,7 +219,7 @@ def process_members(token: str, organization: Organization, existing_teams: List
                                    f"Only the first {MAX_MEMBERS_PER_TEAM} members will be set.")
                     team_members = team_members[:MAX_MEMBERS_PER_TEAM]
                 set_manual_team_members(
-                    token, team_id, team_members, team_name)
+                    token, team_id, team_members, team_name, verify_github_membership)
             else:
                 logger.warning(
                     f"Team '{team_name}' not found in existing teams. Skipping member processing.")
@@ -248,19 +257,24 @@ def main():
         logger.error("Failed to retrieve JWT token. Exiting...")
         return
 
-    organization: Organization = parse_input_file()
+    organization, skip_no_resources, verify_github_membership = parse_input_file()
+    logger.info(
+        f"Running with {skip_no_resources=}, {verify_github_membership=}")
     if not organization:
         logger.error("Failed to parse input file. Exiting...")
         return
 
     assets: List[Asset] = list_assets(jit_token)
 
     existing_teams = get_existing_teams(jit_token)
+    logger.info(
+        f"Found {len(existing_teams)} existing team(s) in the organization: {[team.name for team in existing_teams]}")
     teams_to_delete, created_teams = process_teams(
-        jit_token, organization, assets, existing_teams)
+        jit_token, organization, assets, existing_teams, skip_no_resources)
     existing_teams: List[TeamAttributes] = existing_teams + created_teams
-    desired_teams = get_desired_teams(assets, organization)
-    process_members(jit_token, organization, existing_teams, desired_teams)
+    desired_teams = get_desired_teams(assets, organization, skip_no_resources)
+    process_members(jit_token, organization, existing_teams,
+                    desired_teams, verify_github_membership)
     update_assets(jit_token, assets, organization, existing_teams)
 
     if teams_to_delete:

src/shared/clients/google.py

@@ -0,0 +1,56 @@
+import os
+
+from google.cloud import bigquery
+from google.oauth2 import service_account
+from loguru import logger
+
+from src.shared.models import TeamStructure, Resource, Organization, ResourceType
+
+bigquery_view_name = os.getenv('BIGQUERY_VIEW_NAME')
+credentials_path = os.getenv('GOOGLE_APPLICATION_CREDENTIALS')
+
+
+def get_teams_from_bigquery_view() -> Organization:
+    try:
+        num_repos = 0
+        logger.info(
+            f"Retrieving teams from BigQuery View: {os.getenv('BIGQUERY_VIEW_NAME')}")
+
+        # Explicitly use service account credentials by specifying the private key file.
+        credentials = service_account.Credentials.from_service_account_file(
+            credentials_path)
+        client = bigquery.Client(credentials=credentials)
+
+        teams = {}
+        fields = [
+            "ownership_team_name",
+            "repos",
+            "managers",
+            "manager_usernames",
+            "members",
+            "member_usernames",
+            "slack_alerting_channel"
+        ]
+
+        # Perform a query
+        query = "SELECT {fields} FROM `{view}` WHERE NOT CONTAINS_SUBSTR(ownership_team_name, 'Kaluza')".format(
+            fields=", ".join(fields), view=bigquery_view_name)
+        query_job = client.query(query)  # API request
+        rows = query_job.result()  # Waits for query to finish
+
+        for row in rows:
+            resources = [Resource(type=ResourceType.GithubRepo, name=repo)
+                         for repo in row.repos]
+            num_repos += len(resources)
+            members = list(dict.fromkeys(
+                row.managers + row.members))  # Remove duplicates & keep the same member order
+            teams[row.ownership_team_name] = TeamStructure(
+                name=row.ownership_team_name, members=members, resources=resources,
+                slack_channel=row.slack_alerting_channel)
+        logger.info(
+            f"Retrieved({len(teams.keys())}) teams {list(teams.keys())} with {num_repos} repos"
+            "from Google BigQuery successfully.")
+        return Organization(teams=list(teams.values()))
+    except Exception as e:
+        logger.error(f"Failed to retrieve teams from GitHub: {str(e)}")
+        return Organization(teams=[])

src/shared/clients/jit.py

@@ -53,42 +53,38 @@ def list_assets(token: str) -> List[Asset]:
 
 
 def get_existing_teams(token: str) -> List[TeamAttributes]:
-    def _handle_resoponse(response, existing_teams):
-        response = response.json()
-        data = response['data']
-        existing_teams.extend(data)
-        after = response['metadata']['after']
-        return after
+    headers = {
+        'authorization': f'Bearer {token}',
+    }
 
-    try:
-        # Make a GET request to the asset API
-        url = f"{get_jit_endpoint_base_url()}/teams?limit=100"
+    params = {
+        'limit': '50',
+        'sort_by': 'score',
+        'sort_order': 'desc',
+        'include_members': 'true',
+    }
 
-        headers = get_request_headers(token)
-        response = requests.get(url, headers=headers)
-        existing_teams = []
-        # Check if the request was successful
+    all_teams = []
+    logger.info("Retrieving teams from with pagination.")
+    while True:
+        response = requests.get(
+            f'{get_jit_endpoint_base_url()}/teams', params=params, headers=headers)
         if response.status_code == 200:
-            after = _handle_resoponse(response, existing_teams)
-            while after:
-                response = requests.get(
-                    f"{url}&after={after}", headers=headers)
-                if response.status_code == 200:
-                    after = _handle_resoponse(response, existing_teams)
-                else:
-                    logger.error(
-                        f"Failed to retrieve teams. Status code: {response.status_code}, {response.text}")
-                    return []
+            response_data = response.json()
+            teams = response_data.get('data', [])
+            all_teams.extend(teams)
+            after = response_data.get('metadata', {}).get('after')
+            logger.info(f"Retrieved {len(teams)} teams in page.")
+            if not after:
+                break
 
-            logger.info("Retrieved existing teams successfully.")
-            return [TeamAttributes(**team) for team in existing_teams]
+            params['after'] = after
         else:
             logger.error(
                 f"Failed to retrieve teams. Status code: {response.status_code}, {response.text}")
-            return []
-    except Exception as e:
-        logger.error(f"Failed to retrieve teams: {str(e)}")
-        return []
+            break
+
+    return [TeamAttributes(**team) for team in all_teams]
 
 
 def delete_teams(token, team_names):
@@ -170,12 +166,14 @@ def add_teams_to_asset(token, asset: Asset, teams: List[str]):
 
 
 def _perform_set_manual_team_members(token: str, team_id: str,
-                                     members: List[str], team_name: str) -> Optional[List[str]]:
+                                     members: List[str], team_name: str,
+                                     verify_github_membership: bool) -> Optional[List[str]]:
     try:
         url = f"{get_jit_endpoint_base_url()}/teams/{team_id}/members"
         headers = get_request_headers(token)
         payload = {
-            "members": members
+            "members": members,
+            "verify_github_membership": verify_github_membership
         }
         response = requests.put(url, json=payload, headers=headers)
         if response.status_code == 200:
@@ -197,13 +195,14 @@ def _perform_set_manual_team_members(token: str, team_id: str,
         return None
 
 
-def set_manual_team_members(token: str, team_id: str, members: List[str], team_name: str) -> None:
+def set_manual_team_members(token: str, team_id: str, members: List[str],
+                            team_name: str, verify_github_membership: bool) -> None:
     retry_count = 0
     failed_members = _perform_set_manual_team_members(
-        token, team_id, members, team_name)
+        token, team_id, members, team_name, verify_github_membership)
     while retry_count <= MAX_RETRIES and failed_members:
         failed_members = _perform_set_manual_team_members(
-            token, team_id, members, team_name)
+            token, team_id, members, team_name, verify_github_membership)
         # We send all members, not just the failed ones. Otherwise it would set the list
         # to only the failed members
         retry_count += 1

src/shared/models.py

@@ -31,6 +31,7 @@ class TeamStructure(BaseModel):
     name: str
     members: List[str] = []
     resources: List[Resource] = []
+    slack_channel: Optional[str] = None
 
 
 class Tag(BaseModel):

src/utils/google_bigquery_to_json_file.py

@@ -0,0 +1,22 @@
+import sys
+
+from dotenv import load_dotenv
+from loguru import logger
+
+from src.shared.clients.google import get_teams_from_bigquery_view
+
+# Load environment variables from .env file.
+load_dotenv()
+logger.remove()  # Remove default handler
+logger_format = "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <2}</level> | {message}"
+logger.add(sys.stderr, format=logger_format)
+
+
+def generate_teams_json_file_from_bigquery():
+    teams = get_teams_from_bigquery_view()
+    with open("teams.json", "w") as file:
+        file.write(teams.model_dump_json(indent=2))
+
+
+if __name__ == '__main__':
+    generate_teams_json_file_from_bigquery()