Merge pull request #3 from jitsecurity/sc-19672-create-a-tool-to-generate-teams-from-github-part3

Sc 19672 create a tool to generate teams from GitHub part3

Author: k-shlomi

Makefile

@@ -13,16 +13,18 @@ configure:
 	read -p "Enter JIT API client ID: " client_id; \
 	read -p "Enter JIT API client secret: " client_secret; \
 	read -p "Enter GitHub Personal token (PAT): " github_token; \
+	read -p "Enter comma separated topic wildcards to exclude the creation of teams(example: *dev*, *test*): " topics_to_exclude; \
 	echo "ORGANIZATION_NAME=$$org_name" > .env; \
 	echo "JIT_CLIENT_ID=$$client_id" >> .env; \
 	echo "JIT_CLIENT_SECRET=$$client_secret" >> .env; \
 	echo "GITHUB_API_TOKEN=$$github_token" >> .env
+	echo "TEAM_WILDCARD_TO_EXCLUDE=topics_to_exclude" >> .env
 
 create-teams:
 	source venv-jit/bin/activate && \
 	export PYTHONPATH=$(CURDIR) && \
 	 python src/utils/github_topics_to_json_file.py && \
-	  python src/scripts/create_teams.py --input teams.json
+	  python src/scripts/create_teams.py teams.json
 
 help:
 	@echo "Usage: make [target]"

README.md

@@ -36,6 +36,13 @@ jit-customer-scripts/
 - Python 3.x
 - Git
 
+## Generating API Keys
+
+* To generate Github Personal Access Token(PAT) refer to
+  this [guide](https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens#creating-a-personal-access-token-classic)
+* To generate a Jit API Key Go to Settings -> Users & Permissions -> API Tokens in your
+  Jit [dashboard](https://platform.jit.io).
+
 ## Installation
 
 1. Clone the repository:
@@ -84,23 +91,29 @@ make create-teams
 
 This command is a convenience utility that extracts the teams to generate from Github topics. \
 It runs these commands:
+
 ```bash
 python src/utils/github_topics_to_json_file.py
-python src/scripts/create_teams.py --input teams.json
+python src/scripts/create_teams.py teams.json
 ```
 
-This command will fetch the repository names and topics from the GitHub API and generate the JSON file. And then it will create the teams and update the assets.
+This command will fetch the repository names and topics from the GitHub API and generate the JSON file. And then it will
+create the teams and update the assets.
 
 ### Using External JSON File
 
-You can also provide a JSON file containing team details using the `--input` argument. The JSON file should have the following structure:
+You can also provide a JSON file containing team details using a command line argument directly. The JSON file should
+have the following structure:
 
 ```json
 {
   "teams": [
     {
       "name": "Team 1",
-      "members": ["user1", "user2"],
+      "members": [
+        "user1",
+        "user2"
+      ],
       "resources": [
         {
           "type": "{resource_type}",
@@ -114,7 +127,10 @@ You can also provide a JSON file containing team details using the `--input` arg
     },
     {
       "name": "Team 2",
-      "members": ["user3", "user4"],
+      "members": [
+        "user3",
+        "user4"
+      ],
       "resources": [
         {
           "type": "{resource_type}",
@@ -126,16 +142,27 @@ You can also provide a JSON file containing team details using the `--input` arg
 }
 ```
 
-To use the `--input` argument, run the following command:
+You can run the command like this:
 
 ```shell
-python scripts/create_teams.py --input path/to/teams.json
+python scripts/create_teams.py path/to/teams.json
 ```
 
 Replace `path/to/teams.json` with the actual path to your JSON file.
 
-## Development
+## Excluding Topics
+
+You can exclude certain topics from being considered when creating teams. \
+To exclude topics, you could add them in the `make configure` command or update this env var in
+the `.env` file: `TEAM_WILDCARD_TO_EXCLUDE`.
 
-To override the default Frontegg authentication endpoint, you can set the `FRONTEGG_AUTH_ENDPOINT` environment variable. If the variable is not set, the default value will be used.
+For example, to exclude topics that contain the word "test", you can set the variable as follows:
+
+    TEAM_WILDCARD_TO_EXCLUDE=*test*
+
+This will exclude topics with names like "test", "test123", and "abc-testing".
+
+## Development
 
-To override Jit's API endpoint, you can set the `JIT_API_ENDPOINT` environment variable. If the variable is not set, the default value will be used.
+To override Jit's API endpoint, you can set the `JIT_API_ENDPOINT` environment variable. If the variable is not set, the
+default value will be used.

requirements-test.txt

@@ -2,4 +2,5 @@
 pytest==7.4.0
 pytest-mock==3.11.1
 pytest-cov==4.0.0
+polyfactory==2.7.2
 flake8>=3.9.2

src/scripts/create_teams.py

@@ -7,15 +7,18 @@
 from dotenv import load_dotenv
 from loguru import logger
 from pydantic import ValidationError
-
 from src.shared.clients.jit import get_existing_teams, create_teams, list_assets, add_teams_to_asset, delete_teams, \
     get_jit_jwt_token
 from src.shared.diff_tools import get_different_items_in_lists
-from src.shared.models import Asset, TeamObject, Organization, TeamTemplate
+from src.shared.models import Asset, TeamAttributes, Organization, TeamStructure, ResourceType
 
 # 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: <8}</level> | {message}"
+logger.add(sys.stderr, format=logger_format)
+
 
 def parse_input_file() -> Organization:
     """
@@ -27,39 +30,34 @@ def parse_input_file() -> Organization:
     # Create the argument parser
     parser = argparse.ArgumentParser(description="Retrieve teams and assets")
 
-    # Add the --input argument
-    parser.add_argument("--input", help="Path to a JSON file")
+    # Add the file argument
+    parser.add_argument("file", help="Path to a JSON file")
 
     # Parse the command line arguments
     args = parser.parse_args()
 
-    # Check if the --input argument is provided
-    if args.input:
-        # Check if the file exists and is a JSON file
-        if not os.path.isfile(args.input):
-            logger.error("Error: File does not exist.")
-            sys.exit(1)
-        if not args.input.endswith(".json"):
-            logger.error("Error: File is not a JSON file.")
-            sys.exit(1)
-
-        # Read the JSON file
-        with open(args.input, "r") as file:
-            json_data = file.read()
-
-        # Parse the JSON data
-        try:
-            data = json.loads(json_data)
-            return Organization(teams=[TeamTemplate(**team) for team in data["teams"]])
-        except ValidationError as e:
-            logger.error(f"Failed to validate input file: {e}")
-            sys.exit(1)
-    else:
-        logger.error("No input file provided.")
+    # 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)
+    if not args.file.endswith(".json"):
+        logger.error("Error: File is not a JSON file.")
         sys.exit(1)
 
+    # Read the JSON file
+    with open(args.file, "r") as file:
+        json_data = file.read()
+
+    # Parse the JSON data
+    try:
+        data = json.loads(json_data)
+        return Organization(teams=[TeamStructure(**team) for team in data["teams"]])
+    except (ValidationError, KeyError) as e:
+        logger.error(f"Failed to validate input file: {e}")
+        sys.exit(1)
 
-def update_assets(token, organization):
+
+def update_assets(token, assets: List[Asset], organization):
     """
     Update the assets with the teams specified in the organization.
 
@@ -71,13 +69,23 @@ def update_assets(token, organization):
         None
     """
     logger.info("Updating assets.")
-    assets: List[Asset] = list_assets(token)
+
     asset_to_team_map = get_teams_for_assets(organization)
+    existing_teams: List[str] = [t.name for t in get_existing_teams(token)]
     for asset in assets:
         teams_to_update = asset_to_team_map.get(asset.asset_name, [])
         if teams_to_update:
-            logger.info(f"Adding teams {teams_to_update} to asset {asset.asset_name}")
+            excluded_teams = get_different_items_in_lists(teams_to_update, existing_teams)
+            if excluded_teams:
+                logger.info(f"Excluding team(s) {excluded_teams} for asset '{asset.asset_name}'")
+                teams_to_update = list(set(teams_to_update) - set(excluded_teams))
+            logger.info(f"Syncing team(s) {teams_to_update} to asset '{asset.asset_name}'")
             add_teams_to_asset(token, asset, teams_to_update)
+        else:
+            asset_has_teams_tag = asset.tags and "team" in [t.name for t in asset.tags]
+            if asset_has_teams_tag:
+                logger.info(f"Removing all teams from asset '{asset.asset_name}'")
+                add_teams_to_asset(token, asset, teams_to_update)
 
 
 def get_teams_to_create(topic_names: List[str], existing_team_names: List[str]) -> List[str]:
@@ -108,7 +116,43 @@ 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 process_teams(token, organization):
+def get_desired_teams(assets: List[Asset], organization: Organization) -> 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.
+
+    Returns:
+        List[str]: The names of the desired teams.
+    """
+    desired_teams = []
+    for team in organization.teams:
+        team_resources = []
+        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:
+            desired_teams.append(team.name)
+
+    wildcards_to_exclude = os.getenv("TEAM_WILDCARD_TO_EXCLUDE", "").split(",")
+    final_desired_teams = []
+    for team_name in desired_teams:
+        exclude_team = False
+        for wildcard in wildcards_to_exclude:
+            wildcard = wildcard.strip().strip("*")
+            if wildcard and wildcard in team_name:
+                exclude_team = True
+                break
+        if not exclude_team:
+            final_desired_teams.append(team_name)
+
+    return final_desired_teams
+
+
+def process_teams(token, organization, assets: List[Asset]) -> List[str]:
     """
     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.
@@ -121,13 +165,14 @@ def process_teams(token, organization):
         List[str]: The names of the teams to delete.
     """
     logger.info("Determining required changes in teams.")
-    desired_teams = [t.name for t in organization.teams]
-    existing_teams: List[TeamObject] = get_existing_teams(token)
+
+    desired_teams = get_desired_teams(assets, organization)
+    existing_teams: List[TeamAttributes] = get_existing_teams(token)
     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)
     if teams_to_create:
-        logger.info(f"Creating {len(teams_to_create)} teams: {teams_to_create}")
+        logger.info(f"Creating {len(teams_to_create)} team(s): {teams_to_create}")
         create_teams(token, teams_to_create)
     return teams_to_delete
 
@@ -145,7 +190,7 @@ def get_teams_for_assets(organization: Organization) -> Dict[str, List[str]]:
     asset_to_team_map = {}
     for team in organization.teams:
         for resource in team.resources:
-            if resource.type == "github_repo":
+            if resource.type == ResourceType.GithubRepo:
                 asset_name = resource.name
                 if asset_name in asset_to_team_map:
                     asset_to_team_map[asset_name].append(team.name)
@@ -166,15 +211,17 @@ def main():
         logger.error("Failed to parse input file. Exiting...")
         return
 
-    teams_to_delete = process_teams(jit_token, organization)
+    assets: List[Asset] = list_assets(jit_token)
+
+    teams_to_delete = process_teams(jit_token, organization, assets)
 
-    update_assets(jit_token, organization)
+    update_assets(jit_token, assets, organization)
 
     if teams_to_delete:
-        logger.info(f"Deleting {len(teams_to_delete)} teams: {teams_to_delete}")
+        logger.info(f"Deleting {len(teams_to_delete)} team(s): {teams_to_delete}")
         delete_teams(jit_token, teams_to_delete)
+    logger.info("Successfully completed teams sync.")
 
 
 if __name__ == '__main__':
-    logger.add("app.log", rotation="5 MB", level="INFO")
     main()

src/shared/clients/github.py

@@ -1,26 +1,19 @@
 import os
 
-from dotenv import load_dotenv
 from github import Github
 from loguru import logger
 
-from src.shared.models import TeamTemplate, Resource, Organization
-
-# Load environment variables from .env file.
-load_dotenv(".env")
-
-ORGANIZATION_NAME = os.getenv("ORGANIZATION_NAME")
-GITHUB_TOKEN = os.getenv("GITHUB_API_TOKEN")
+from src.shared.models import TeamStructure, Resource, Organization, ResourceType
 
 
 def get_teams_from_github_topics() -> Organization:
     try:
-        logger.info(f"Trying to communicate with Github to get information from Org: {ORGANIZATION_NAME}")
+        logger.info(f"Trying to communicate with Github to get information from Org: {os.getenv('ORGANIZATION_NAME')}")
         # Create a GitHub instance using the token
-        github = Github(GITHUB_TOKEN)
+        github = Github(os.getenv("GITHUB_API_TOKEN"))
 
         # Get the organization
-        organization = github.get_organization(ORGANIZATION_NAME)
+        organization = github.get_organization(os.getenv('ORGANIZATION_NAME'))
 
         # Dictionary to store team templates
         teams = {}
@@ -40,16 +33,16 @@ def get_teams_from_github_topics() -> Organization:
                 # Check if the topic already exists in the teams dictionary
                 if topic in teams:
                     # Add the repository to the existing team
-                    teams[topic].resources.append(Resource(type="github_repo", name=repo_name))
+                    teams[topic].resources.append(Resource(type=ResourceType.GithubRepo, name=repo_name))
                 else:
                     # Create a new team template for the topic
-                    team_template = TeamTemplate(name=topic, members=[],
-                                                 resources=[Resource(type="github_repo", name=repo_name)])
+                    team_template = TeamStructure(name=topic, members=[],
+                                                  resources=[Resource(type=ResourceType.GithubRepo, name=repo_name)])
 
                     # Add the team template to the teams dictionary
                     teams[topic] = team_template
 
-        logger.info(f"Retrieved ({len(teams.keys())}) teams {list(teams.keys())} from GitHub successfully.")
+        logger.info(f"Retrieved ({len(teams.keys())}) team(s) {list(teams.keys())} from GitHub successfully.")
         return Organization(teams=list(teams.values()))
     except Exception as e:
         logger.error(f"Failed to retrieve teams from GitHub: {str(e)}")