Skip to content

Commit 7e7d641

Browse files
committed
Add asserter to check for organization plan name
1 parent 27cb27a commit 7e7d641

File tree

1 file changed

+21
-5
lines changed

1 file changed

+21
-5
lines changed

socketsync/core.py

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -120,8 +120,10 @@ def set_org_vars() -> None:
120120
:return:
121121
"""
122122
log.debug("Getting Organization Configuration")
123-
global org_id, org_slug, full_scan_path, repository_path, security_policy
124-
org_id, org_slug = Core.get_org_id_slug()
123+
global org_id, org_slug, org_plan, full_scan_path, repository_path, security_policy
124+
org_id, org_slug, org_plan = Core.get_org_id_slug_plan()
125+
# Ensure the organization is on the Enterprise plan
126+
Core.assert_enterprise_plan()
125127
base_path = f"orgs/{org_slug}"
126128
full_scan_path = f"{base_path}/full-scans"
127129
repository_path = f"{base_path}/repos"
@@ -148,20 +150,34 @@ def set_timeout(request_timeout: int):
148150
socketdev.set_timeout(timeout)
149151

150152
@staticmethod
151-
def get_org_id_slug() -> (str, str):
153+
def get_org_id_slug_plan() -> (str, str, str):
152154
"""
153-
Gets the Org ID and Org Slug for the API Token
155+
Gets the Org ID, Org Slug, and Org Plan for the API Token
154156
:return:
155157
"""
156158
organizations = socket.org.get()
157159
orgs = organizations.get("organizations")
158160
new_org_id = None
159161
new_org_slug = None
162+
new_org_plan = None
160163
if orgs is not None and len(orgs) == 1:
161164
for key in orgs:
162165
new_org_id = key
163166
new_org_slug = orgs[key].get("slug")
164-
return new_org_id, new_org_slug
167+
new_org_plan = orgs[key].get("plan")
168+
return new_org_id, new_org_slug, new_org_plan
169+
170+
@staticmethod
171+
def assert_enterprise_plan() -> None:
172+
"""
173+
Validate that the current organization is on the Enterprise plan. If not, fail fast.
174+
:param organization_id: The organization id to check
175+
:return: None
176+
:raises Exception: if the organization is not on the Enterprise plan
177+
"""
178+
is_enterprise = "enterprise" in org_plan.lower()
179+
if not is_enterprise:
180+
raise Exception("This script requires an Enterprise plan organization")
165181

166182
@staticmethod
167183
def get_security_policy() -> dict:

0 commit comments

Comments
 (0)