diff --git a/src/crowdstrike/foundry/function/falconpy.py b/src/crowdstrike/foundry/function/falconpy.py index 452b980..86812dd 100644 --- a/src/crowdstrike/foundry/function/falconpy.py +++ b/src/crowdstrike/foundry/function/falconpy.py @@ -29,5 +29,9 @@ def falcon_client(client_class: Type) -> ServiceClass: cloud = cloud.lower().replace('-', '').strip() if cloud == '': cloud = _cloud_default + + # set cloud on the request object that is in the context + req.cloud = cloud + client = client_class(access_token=access_token, base_url=cloud) return client diff --git a/src/crowdstrike/foundry/function/model.py b/src/crowdstrike/foundry/function/model.py index 549a0c3..4d769b7 100644 --- a/src/crowdstrike/foundry/function/model.py +++ b/src/crowdstrike/foundry/function/model.py @@ -22,6 +22,7 @@ class Request: method: str = field(default='') params: RequestParams = field(default_factory=lambda: RequestParams()) url: str = field(default='') + cloud: str = field(default='') @dataclass diff --git a/tests/crowdstrike/foundry/function/test_falconpy.py b/tests/crowdstrike/foundry/function/test_falconpy.py index 3b43222..abbe168 100644 --- a/tests/crowdstrike/foundry/function/test_falconpy.py +++ b/tests/crowdstrike/foundry/function/test_falconpy.py @@ -2,8 +2,8 @@ import unittest from unittest.mock import patch from crowdstrike.foundry.function.context import ctx_request -from crowdstrike.foundry.function.falconpy import falcon_client -from crowdstrike.foundry.function.model import Request +from crowdstrike.foundry.function.falconpy import falcon_client +from crowdstrike.foundry.function.model import Request from falconpy import Hosts if __name__ == '__main__': @@ -54,3 +54,14 @@ def test_request_with_access_token_and_non_default_cloud_returns_prepared_client self.assertIsInstance(client, Hosts) self.assertEqual('Bearer foo', client.headers.get('Authorization')) self.assertEqual(t['expected'], client.base_url) + + def test_request_inserts_cloud_into_request(self): + with patch.dict(os.environ, {'CS_CLOUD': 'us-gov-1'}, clear=True): + ctx_request.set(Request(access_token='foo')) + client = falcon_client(Hosts) + + self.assertIsInstance(client, Hosts) + self.assertEqual('Bearer foo', client.headers.get('Authorization')) + self.assertEqual('https://api.laggar.gcw.crowdstrike.com', client.base_url) + r = ctx_request.get() + self.assertEqual('usgov1', r.cloud)