@@ -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