Merge pull request #20 from snyk-labs/feat/mt-support

feat: Multitenant EU and AU support
snyk-labs · Sep 10, 2023 · 5c69a5c · 5c69a5c
2 parents 3e049cd + 5007c65
commit 5c69a5c
Show file tree

Hide file tree

Showing 9 changed files with 293 additions and 124 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "snyk-tags"
-version = "2.1.0"
+version = "2.2.0"
 description = "CLI tool designed to manage tags and attributes at scale"
 authors = ["EricFernandezSnyk <[email protected]>"]
 license = "MIT"

diff --git a/snyk_tags/__init__.py b/snyk_tags/__init__.py
@@ -1,7 +1,7 @@
 # snyk_tags/__init__.py
 
 __app_name__ = "snyk_tags"
-__version__ = "2.1.0"
+__version__ = "2.2.0"
 
 from logging import ERROR
 from sre_constants import SUCCESS

diff --git a/snyk_tags/attribute.py b/snyk_tags/attribute.py
@@ -20,11 +20,14 @@
 
 
 # Reach to the API and generate tokens
-def create_client(token: str) -> httpx.Client:
-    return httpx.Client(
-        base_url="https://snyk.io/api/v1",
-        headers={"Authorization": f"token {token}", "Content-Type": "application/json"},
+def create_client(token: str, tenant: str) -> httpx.Client:
+    base_url = (
+        f"https://api.{tenant}.snyk.io/v1"
+        if tenant in ["eu", "au"]
+        else "https://api.snyk.io/v1"
     )
+    headers = {"Authorization": f"token {token}"}
+    return httpx.Client(base_url=base_url, headers=headers)
 
 
 # Apply attributes to a specific project
@@ -66,19 +69,20 @@ def apply_attributes_to_project(
     criticality = typer.style(criticality, bold=True, fg=typer.colors.MAGENTA)
     environment = typer.style(environment, bold=True, fg=typer.colors.MAGENTA)
     lifecycle = typer.style(lifecycle, bold=True, fg=typer.colors.MAGENTA)
+
     if req.status_code == 200:
         logging.info(
             f"Successfully added {criticality},{environment},{lifecycle} attributes to Project: {project_name}."
         )
-    if req.status_code == 422:
+    elif req.status_code == 422:
         logging.warning(
             f"Data {attribute_data} cannot be processed, make sure you have written the correct values (refer to help or Readme) and that they are in low caps. Error message: {req.json()}."
         )
-    if req.status_code == 404:
+    elif req.status_code == 404:
         logging.error(
             f"Project not found, likely a READ-ONLY project. Project: {project_name}. Error message: {req.json()}."
         )
-    if req.status_code == 500:
+    elif req.status_code == 500:
         logging.error(
             f"Error message: {req.json()}. Please contact [email protected]."
         )
@@ -93,25 +97,33 @@ def apply_attributes_to_projects(
     criticality: list,
     environment: list,
     lifecycle: list,
+    tenant: str,
 ) -> None:
-    with create_client(token=token) as client:
+    with create_client(token=token, tenant=tenant) as client:
         for org_id in org_ids:
-            client_v3 = SnykClient(token=token)
-            projects = client_v3.organizations.get(org_id).projects.all()
+            base_url = (
+                f"https://api.{tenant}.snyk.io/rest"
+                if tenant in ["eu", "au"]
+                else "https://api.snyk.io/rest"
+            )
+            client_v3 = SnykClient(
+                token=token, url=base_url, version="2023-08-31~experimental"
+            )
+            projects = client_v3.get(f"/orgs/{org_id}/projects").json()
 
             badname = 0
             rightname = 0
 
-            for project in projects:
-                if project.name.startswith(name):
+            for project in projects["data"]:
+                if project["attributes"]["name"].startswith(name):
                     apply_attributes_to_project(
                         client=client,
                         org_id=org_id,
-                        project_id=project.id,
+                        project_id=project["id"],
                         criticality=criticality,
                         environment=environment,
                         lifecycle=lifecycle,
-                        project_name=project.name,
+                        project_name=project["attributes"]["name"],
                     )
                     rightname = 1
                 else:

diff --git a/snyk_tags/collection.py b/snyk_tags/collection.py
@@ -24,10 +24,14 @@
 
 
 # Reach to the API and generate tokens
-def create_client(token: str) -> httpx.Client:
-    return httpx.Client(
-        base_url="https://snyk.io/api/v1", headers={"Authorization": f"token {token}"}
+def create_client(token: str, tenant: str) -> httpx.Client:
+    base_url = (
+        f"https://api.{tenant}.snyk.io/v1"
+        if tenant in ["eu", "au"]
+        else "https://api.snyk.io/v1"
     )
+    headers = {"Authorization": f"token {token}"}
+    return httpx.Client(base_url=base_url, headers=headers)
 
 
 # Apply tags to a specific project
@@ -50,11 +54,11 @@ def apply_tag_to_project(
 
     if req.status_code == 200:
         logging.info(f"Successfully added {tag_data} tags to Project: {project_name}.")
-    if req.status_code == 422:
+    elif req.status_code == 422:
         logging.warning(
             f"Tag {key}:{tag} is already applied for Project: {project_name}."
         )
-    if req.status_code == 404:
+    elif req.status_code == 404:
         logging.error(
             f"Project not found, likely a READ-ONLY project. Project: {project_name}. Error message: {req.json()}."
         )
@@ -63,28 +67,35 @@ def apply_tag_to_project(
 
 # Tagging loop
 def apply_tags_to_projects(
-    token: str, org_ids: list, name: str, tag: str, key: str
+    token: str, org_ids: list, name: str, tag: str, key: str, tenant: str
 ) -> None:
-    with create_client(token=token) as client:
+    with create_client(token=token, tenant=tenant) as client:
         for org_id in org_ids:
-            client_v3 = SnykClient(token=token)
-            projects = client_v3.organizations.get(org_id).projects.all()
+            base_url = (
+                f"https://api.{tenant}.snyk.io/rest"
+                if tenant in ["eu", "au"]
+                else "https://api.snyk.io/rest"
+            )
+            client_v3 = SnykClient(
+                token=token, url=base_url, version="2023-08-31~experimental"
+            )
+            projects = client_v3.get(f"/orgs/{org_id}/projects").json()
 
             badname = 0
             rightname = 0
-            for project in projects:
+            for project in projects["data"]:
                 if (
-                    project.name == name
-                    or project.name.startswith(name + "(")
-                    or project.name.startswith(name + ":")
+                    project["attributes"]["name"] == name
+                    or project["attributes"]["name"](name + "(")
+                    or project["attributes"]["name"](name + ":")
                 ):
                     apply_tag_to_project(
                         client=client,
                         org_id=org_id,
-                        project_id=project.id,
+                        project_id=project["id"],
                         tag=tag,
                         key=key,
-                        project_name=project.name,
+                        project_name=project["attributes"]["name"],
                     )
                     rightname = 1
                 else:
@@ -124,6 +135,10 @@ def tag(
         ...,  # Default value of comamand
         help=f"Name of the target, for example {repoexample}",
     ),
+    tenant: str = typer.Option(
+        "",  # Default value of comamand
+        help=f"Defaults to US tenant, add 'eu' or 'au' to use EU or AU tenant, use --tenant to change tenant.",
+    ),
     tagKey: str = typer.Option(
         ..., help="Tag key: identifier of the tag"  # Default value of comamand
     ),
@@ -136,7 +151,7 @@ def tag(
         bold=True,
         fg=typer.colors.MAGENTA,
     )
-    apply_tags_to_projects(snyktkn, [org_id], target, tagValue, tagKey)
+    apply_tags_to_projects(snyktkn, [org_id], target, tagValue, tagKey, tenant=tenant)
 
 
 # Collection command to apply the attributes to the collection
@@ -165,12 +180,22 @@ def attributes(
     lifecycle: str = typer.Option(
         "", help=f"Lifecycle attribute: {life}"  # Default value of comamand
     ),
+    tenant: str = typer.Option(
+        "",  # Default value of comamand
+        help=f"Defaults to US tenant, add 'eu' or 'au' to use EU or AU tenant, use --tenant to change tenant.",
+    ),
 ):
     typer.secho(
         f"\nAdding the attributes {criticality}, {environment} and {lifecycle} to projects within {target} for easy filtering via the UI",
         bold=True,
         fg=typer.colors.MAGENTA,
     )
     attribute.apply_attributes_to_projects(
-        snyktkn, [org_id], target, [criticality], [environment], [lifecycle]
+        snyktkn,
+        [org_id],
+        target,
+        [criticality],
+        [environment],
+        [lifecycle],
+        tenant=tenant,
     )
diff --git a/snyk_tags/files.py b/snyk_tags/files.py
@@ -31,6 +31,10 @@ def target_tag(
     snyktkn: str = typer.Option(
         ..., help="Snyk API token with org admin access", envvar=["SNYK_TOKEN"]
     ),
+    tenant: str = typer.Option(
+        "",  # Default value of comamand
+        help=f"Defaults to US tenant, add 'eu' or 'au' to use EU or AU tenant, use --tenant to change tenant.",
+    ),
 ):
     for path in file:
         if path.is_file():
@@ -47,7 +51,7 @@ def target_tag(
                         bold=True,
                     )
                     collection.apply_tags_to_projects(
-                        snyktkn, [org_id], target, value, key
+                        snyktkn, [org_id], target, value, key, tenant
                     )
                 openfile.close()
             elif ".json" in openfile.name:
@@ -62,7 +66,7 @@ def target_tag(
                         bold=True,
                     )
                     collection.apply_tags_to_projects(
-                        snyktkn, [org_id], target, value, key
+                        snyktkn, [org_id], target, value, key, tenant
                     )
                 openfile.close()
             else:
@@ -84,6 +88,10 @@ def target_attributes(
     snyktkn: str = typer.Option(
         ..., help="Snyk API token with org admin access", envvar=["SNYK_TOKEN"]
     ),
+    tenant: str = typer.Option(
+        "",  # Default value of comamand
+        help=f"Defaults to US tenant, add 'eu' or 'au' to use EU or AU tenant, use --tenant to change tenant.",
+    ),
 ):
     for path in file:
         if path.is_file():
@@ -108,6 +116,7 @@ def target_attributes(
                         [criticality],
                         [environment],
                         [lifecycle],
+                        tenant,
                     )
                 openfile.close()
             elif ".json" in openfile.name:
@@ -130,6 +139,7 @@ def target_attributes(
                         [criticality],
                         [environment],
                         [lifecycle],
+                        tenant,
                     )
                 openfile.close()
             else: