use jit endpoint for authentication

Author: k-shlomi

src/scripts/create_teams.py

@@ -8,8 +8,8 @@
 from loguru import logger
 from pydantic import ValidationError
 
-from src.shared.clients.frontegg import get_jit_jwt_token
-from src.shared.clients.jit import get_existing_teams, create_teams, list_assets, add_teams_to_asset, delete_teams
+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
 

src/shared/clients/frontegg.py

@@ -1,5 +1,6 @@
 import os
 from typing import List
+from typing import Optional
 
 import requests
 from dotenv import load_dotenv
@@ -12,6 +13,28 @@
 load_dotenv(".env")
 JIT_API_ENDPOINT = os.getenv("JIT_API_ENDPOINT", JIT_DEFAULT_API_ENDPOINT)
 
+JIT_CLIENT_SECRET = os.getenv("JIT_CLIENT_SECRET")
+JIT_CLIENT_ID = os.getenv("JIT_CLIENT_ID")
+
+
+def get_jit_jwt_token() -> Optional[str]:
+    payload = {
+        "clientId": JIT_CLIENT_ID,
+        "secret": JIT_CLIENT_SECRET
+    }
+    headers = {
+        "accept": "application/json",
+        "content-type": "application/json"
+    }
+
+    response = requests.post(f"{JIT_API_ENDPOINT}/authentication/login", json=payload, headers=headers)
+
+    if response.status_code == 200:
+        return response.json().get('accessToken')
+    else:
+        logger.error(f"Failed to retrieve JWT token. Status code: {response.status_code}")
+        return None
+
 
 def list_assets(token: str) -> List[Asset]:
     try:
@@ -127,7 +150,7 @@ def get_request_headers(token):
 
 def add_teams_to_asset(token, asset: Asset, teams: List[str]):
     try:
-        url = f"{JIT_API_ENDPOINT}/assets/{asset.asset_id}"
+        url = f"{JIT_API_ENDPOINT}/asset/asset/{asset.asset_id}"
         headers = get_request_headers(token)
         payload = {
             "teams": teams

src/shared/clients/jit.py

@@ -1,4 +1,3 @@
 MANUAL_TEAM_SOURCE = "manual"
 REPO = 'repo'
-FRONTEGG_DEFAULT_AUTH_URL = "https://jit.frontegg.com/identity/resources/auth/v1/api-token"
 JIT_DEFAULT_API_ENDPOINT = "https://api.jit.io"

src/shared/consts.py

@@ -1,8 +1,8 @@
 import pytest
 
-from src.shared.clients.frontegg import get_jit_jwt_token, FRONTEGG_AUTH_URL
 from src.shared.clients.github import get_teams_from_github_topics
-from src.shared.clients.jit import list_assets, get_existing_teams, create_teams, add_teams_to_asset, delete_teams
+from src.shared.clients.jit import list_assets, get_existing_teams, create_teams, add_teams_to_asset, delete_teams, \
+    get_jit_jwt_token, JIT_API_ENDPOINT
 from src.shared.models import TeamObject, Asset, Organization, TeamTemplate, Resource
 
 
@@ -90,7 +90,7 @@ def test_get_jwt_token(status_code, expected_result, mocker):
     token = get_jit_jwt_token()
 
     requests_post_mock.assert_called_once_with(
-        FRONTEGG_AUTH_URL,
+        f"{JIT_API_ENDPOINT}/authentication/login",
         json={"clientId": None, "secret": None},
         headers={"accept": "application/json", "content-type": "application/json"}
     )
@@ -203,7 +203,7 @@ def test_delete_teams(mocker, status_code, existing_team_names, input_team_names
                       expected_warning):
     mock_existing_teams = [
         TeamObject(tenant_id=f"tenant{i + 1}", id=str(i + 1), created_at=f"date{i + 1}", modified_at=f"date{i + 2}",
-                 name=team_name)
+                   name=team_name)
         for i, team_name in enumerate(existing_team_names)
     ]