Skip to content

Commit 44b97ba

Browse files
idanchAquaidanchAqua
committed
SLK-103546 - Add Token Authentication Fallback Support
- Add X-Tokens-Signature header to Auto-Connect backend API calls - Implement automatic Bearer token fallback for direct CSPM calls on 401/403 errors - Add get_bearer_token() and cspm_request_with_fallback() helper functions - Refactor to use cspm_base_url and path parameters, remove duplicate URL parsing - Add logging for HTTP requests, responses, and fallback operations - Update trigger-aws.py, create_cspm_key.py, and generate_external_id.py
1 parent 5699472 commit 44b97ba

File tree

3 files changed

+218
-9
lines changed

3 files changed

+218
-9
lines changed

‎modules/single/modules/lambda/functions/create_cspm_key.py‎

Lines changed: 70 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,22 +44,90 @@ def http_request(url, headers, method, body=None):
4444
if body is None:
4545
body = {}
4646

47+
print(f"HTTP request: {method} {url}")
48+
4749
http = urllib3.PoolManager(cert_reqs='CERT_NONE')
4850

4951
try:
5052
response = http.request(method, url, body=body, headers=headers)
53+
response_data = response.data.decode('utf-8') if response.data else ''
54+
print(f"HTTP response: {response.status} {response.reason} - {response_data}")
5155
return response
5256
except Exception as e:
5357
print('Failed to send http request; {}'.format(e))
5458
return None
5559

5660

61+
def get_bearer_token(cspm_base_url, api_key, aqua_secret, tstmp):
62+
"""Obtain Bearer token from CSPM /v2/tokens endpoint"""
63+
path = "/v2/tokens"
64+
method = "POST"
65+
body = '{"validity":1,"allowed_endpoints":["ANY"]}'
66+
tokens_url = cspm_base_url + path
67+
68+
print("Token fallback: Calling POST /v2/tokens to obtain Bearer token")
69+
70+
tokens_sig = get_signature(aqua_secret, tstmp, path, method, body)
71+
headers = {
72+
"X-API-Key": api_key,
73+
"X-Signature": tokens_sig,
74+
"X-Timestamp": tstmp,
75+
"Content-Type": "application/json"
76+
}
77+
78+
response = http_request(tokens_url, headers, method, body)
79+
if response is None:
80+
raise Exception("Failed to get Bearer token: HTTP request failed")
81+
82+
if response.status not in [200, 201]:
83+
raise Exception(f"Failed to get Bearer token: {response.data.decode('utf-8')}")
84+
85+
json_object = json.loads(response.data.decode('utf-8'))
86+
if json_object.get('status') != 200:
87+
error_msg = json_object.get('message', 'Unknown error')
88+
raise Exception(f"Tokens API failed: {error_msg}")
89+
90+
return json_object['data']
91+
92+
93+
def cspm_request_with_fallback(cspm_base_url, path, headers, method, body, api_key, aqua_secret, tstmp):
94+
"""Make CSPM request with automatic token authentication fallback"""
95+
url = cspm_base_url + path
96+
response = http_request(url, headers, method, body if body else '')
97+
98+
if response is None:
99+
raise ValueError("HTTP request failed")
100+
101+
# Attempt fallback for 401/403 errors
102+
if response.status in [401, 403]:
103+
print(f"Token fallback: API key authentication failed with status {response.status}, attempting Bearer token fallback")
104+
try:
105+
bearer_token = get_bearer_token(cspm_base_url, api_key, aqua_secret, tstmp)
106+
107+
fallback_headers = {
108+
"Authorization": f"Bearer {bearer_token}",
109+
"X-Timestamp": tstmp,
110+
"Content-Type": "application/json"
111+
}
112+
113+
response = http_request(url, fallback_headers, method, body if body else '')
114+
if response and response.status in [200, 201]:
115+
print("Token fallback: Bearer token authentication succeeded")
116+
elif response:
117+
print(f"Token fallback: Bearer token authentication failed with status {response.status}")
118+
except Exception as e:
119+
print(f"Token fallback failed: {e}")
120+
# Return original response if fallback fails
121+
122+
return response
123+
124+
57125
def get_cspm_key_id(aqua_api_key, aqua_secret, cspm_url, role_arn):
58126
tstmp = str(int(time.time() * 1000))
59127
sig = get_signature(aqua_secret, tstmp, "/v2/keys", "GET", '')
60128
headers = {"X-API-Key": aqua_api_key, "X-Signature": sig, "X-Timestamp": tstmp}
61129

62-
response = http_request(cspm_url + "/v2/keys", headers, "GET")
130+
response = cspm_request_with_fallback(cspm_url, "/v2/keys", headers, "GET", '', aqua_api_key, aqua_secret, tstmp)
63131
json_object = json.loads(response.data)
64132
if response.status not in (200, 201):
65133
raise ValueError(f"Failed to get cspm key id for {role_arn}: {response.message}")
@@ -93,7 +161,7 @@ def create_cspm_key(cspm_url, aqua_api_key, aqua_secret, role_arn, external_id,
93161
"X-Timestamp": tstmp
94162
}
95163

96-
response = http_request(cspm_url + '/v2/keys', headers, "POST", jsonbody)
164+
response = cspm_request_with_fallback(cspm_url, '/v2/keys', headers, "POST", jsonbody, aqua_api_key, aqua_secret, tstmp)
97165
if response.status not in (200, 201):
98166
raise Exception("Failed to create cspm key id", response.data.decode("utf-8"))
99167

‎modules/single/modules/lambda/functions/generate_external_id.py‎

Lines changed: 74 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,81 @@ def http_request(url, headers, method, body=None):
3333
if body is None:
3434
body = {}
3535

36+
print(f"HTTP request: {method} {url}")
37+
3638
http = urllib3.PoolManager(cert_reqs='CERT_REQUIRED')
3739

3840
try:
3941
response = http.request(method, url, body=body, headers=headers)
42+
response_data = response.data.decode('utf-8') if response.data else ''
43+
print(f"HTTP response: {response.status} {response.reason} - {response_data}")
44+
return response
45+
except Exception as e:
46+
print('Failed to send http request; {}'.format(e))
47+
return None
48+
49+
50+
def get_bearer_token(cspm_base_url, api_key, aqua_secret, tstmp):
51+
"""Obtain Bearer token from CSPM /v2/tokens endpoint"""
52+
path = "/v2/tokens"
53+
method = "POST"
54+
body = '{"validity":1,"allowed_endpoints":["ANY"]}'
55+
tokens_url = cspm_base_url + path
56+
57+
print("Token fallback: Calling POST /v2/tokens to obtain Bearer token")
58+
59+
tokens_sig = get_signature(aqua_secret, tstmp, path, method, body)
60+
headers = {
61+
"X-API-Key": api_key,
62+
"X-Signature": tokens_sig,
63+
"X-Timestamp": tstmp,
64+
"Content-Type": "application/json"
65+
}
66+
67+
response = http_request(tokens_url, headers, method, body)
68+
69+
if response.status not in [200, 201]:
70+
raise Exception(f"Failed to get Bearer token: {response.data.decode('utf-8')}")
4071

72+
json_object = json.loads(response.data.decode('utf-8'))
73+
if json_object.get('status') != 200:
74+
error_msg = json_object.get('message', 'Unknown error')
75+
raise Exception(f"Tokens API failed: {error_msg}")
76+
77+
return json_object['data']
78+
79+
80+
def cspm_request_with_fallback(cspm_base_url, path, headers, method, body, api_key, aqua_secret, tstmp):
81+
"""Make CSPM request with automatic token authentication fallback"""
82+
url = cspm_base_url + path
83+
response = http_request(url, headers, method, body if body else '')
84+
85+
if response is None:
86+
raise ValueError("HTTP request failed")
87+
88+
# Attempt fallback for 401/403 errors
89+
if response.status in [401, 403]:
90+
print(f"Token fallback: API key authentication failed with status {response.status}, attempting Bearer token fallback")
91+
try:
92+
bearer_token = get_bearer_token(cspm_base_url, api_key, aqua_secret, tstmp)
93+
94+
fallback_headers = {
95+
"Authorization": f"Bearer {bearer_token}",
96+
"X-Timestamp": tstmp,
97+
"Content-Type": "application/json"
98+
}
99+
100+
response = http_request(url, fallback_headers, method, body if body else '')
101+
if response and response.status in [200, 201]:
102+
print("Token fallback: Bearer token authentication succeeded")
103+
elif response:
104+
print(f"Token fallback: Bearer token authentication failed with status {response.status}")
105+
except Exception as e:
106+
print(f"Token fallback failed: {e}")
107+
# Continue with original response if fallback fails
108+
109+
# Parse response to match existing http_request behavior
110+
try:
41111
data = json.loads(response.data.decode('utf-8'))
42112
except Exception as e:
43113
print("warning: {}".format(e))
@@ -47,15 +117,15 @@ def http_request(url, headers, method, body=None):
47117

48118

49119
def generate_external_id(cspm_url, ac_url, aqua_api_key, aqua_secret, aws_account_id):
50-
u = cspm_url + '/v2/generatedids'
51-
print('api url: {}'.format(u))
120+
path = '/v2/generatedids'
121+
print('api url: {}'.format(cspm_url + path))
52122

53123
tstmp = str(int(time.time() * 1000))
54124
method = "POST"
55-
sig = get_signature(aqua_secret, tstmp, '/v2/generatedids', method, '')
125+
sig = get_signature(aqua_secret, tstmp, path, method, '')
56126
headers = {"X-API-Key": aqua_api_key, "X-Signature": sig, "X-Timestamp": tstmp}
57127

58-
response = http_request(u, headers, method)
128+
response = cspm_request_with_fallback(cspm_url, path, headers, method, '', aqua_api_key, aqua_secret, tstmp)
59129
if response.get('status', 0) != 200 and response.get('status', 0) != 201 or not response.get('data'):
60130
raise Exception("failed to generate external id; {}".format(response.get('message', 'Internal server error')))
61131

‎modules/single/modules/trigger/trigger-aws.py‎

Lines changed: 74 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@ def http_request(url, headers, method, body=None):
4949
else:
5050
raise ValueError("Unsupported URL format")
5151

52+
print(f"HTTP request: {method} {url}")
53+
5254
try:
5355
conn = http.client.HTTPSConnection(hostname, context=ssl._create_unverified_context())
5456
conn.request(method, path, body=body, headers=headers)
@@ -58,6 +60,8 @@ def http_request(url, headers, method, body=None):
5860

5961
conn.close()
6062

63+
print(f"HTTP response: {response.status} {response.reason} - {response_data}")
64+
6165
return {
6266
"status": response.status,
6367
"reason": response.reason,
@@ -68,6 +72,70 @@ def http_request(url, headers, method, body=None):
6872
return None
6973

7074

75+
def get_bearer_token(cspm_base_url, api_key, aqua_secret, tstmp):
76+
"""Obtain Bearer token from CSPM /v2/tokens endpoint"""
77+
path = "/v2/tokens"
78+
method = "POST"
79+
body = '{"validity":1,"allowed_endpoints":["ANY"]}'
80+
tokens_url = cspm_base_url + path
81+
82+
print("Token fallback: Calling POST /v2/tokens to obtain Bearer token")
83+
84+
tokens_sig = get_signature(aqua_secret, tstmp, path, method, body)
85+
headers = {
86+
"X-API-Key": api_key,
87+
"X-Signature": tokens_sig,
88+
"X-Timestamp": tstmp,
89+
"Content-Type": "application/json"
90+
}
91+
92+
response = http_request(tokens_url, headers, method, body)
93+
if response is None:
94+
raise Exception("Failed to get Bearer token: HTTP request failed")
95+
96+
if response["status"] not in [200, 201]:
97+
raise Exception(f"Failed to get Bearer token: {response['data']}")
98+
99+
json_object = json.loads(response["data"])
100+
if json_object.get('status') != 200:
101+
error_msg = json_object.get('message', 'Unknown error')
102+
raise Exception(f"Tokens API failed: {error_msg}")
103+
104+
return json_object['data']
105+
106+
107+
def cspm_request_with_fallback(cspm_base_url, path, headers, method, body, api_key, aqua_secret, tstmp):
108+
"""Make CSPM request with automatic token authentication fallback"""
109+
url = cspm_base_url + path
110+
response = http_request(url, headers, method, body)
111+
112+
if response is None:
113+
raise ValueError("HTTP request failed")
114+
115+
# Attempt fallback for 401/403 errors
116+
if response["status"] in [401, 403]:
117+
print(f"Token fallback: API key authentication failed with status {response['status']}, attempting Bearer token fallback")
118+
try:
119+
bearer_token = get_bearer_token(cspm_base_url, api_key, aqua_secret, tstmp)
120+
121+
fallback_headers = {
122+
"Authorization": f"Bearer {bearer_token}",
123+
"X-Timestamp": tstmp,
124+
"Content-Type": "application/json"
125+
}
126+
127+
response = http_request(url, fallback_headers, method, body)
128+
if response["status"] in [200, 201]:
129+
print("Token fallback: Bearer token authentication succeeded")
130+
else:
131+
print(f"Token fallback: Bearer token authentication failed with status {response['status']}")
132+
except Exception as e:
133+
print(f"Token fallback failed: {e}")
134+
# Return original response if fallback fails
135+
136+
return response
137+
138+
71139
def get_cspm_key_id(aqua_api_key, aqua_secret, cspm_url, role_arn):
72140
"""Fetch the CSPM key ID for the given IAM role ARN"""
73141

@@ -78,7 +146,7 @@ def get_cspm_key_id(aqua_api_key, aqua_secret, cspm_url, role_arn):
78146
"X-Timestamp": tstmp
79147
}
80148

81-
response = http_request(cspm_url + "/v2/keys", headers, "GET")
149+
response = cspm_request_with_fallback(cspm_url, "/v2/keys", headers, "GET", '', aqua_api_key, aqua_secret, tstmp)
82150

83151
if response is None:
84152
raise ValueError(f"HTTP request failed while getting CSPM key ID for {role_arn}")
@@ -130,10 +198,12 @@ def trigger_discovery():
130198
)
131199

132200
cspm_sig = get_signature(aqua_secret, tstmp, "/v2/keys", "POST", body_cspm)
201+
tokens_signature = get_signature(aqua_secret, tstmp, "/v2/tokens", "POST", '{"validity":1,"allowed_endpoints":["ANY"]}')
133202
headers = {
134203
"X-API-Key": aqua_api_key,
135204
"X-Authenticate-Api-Key-Signature": sig,
136205
"X-Register-New-Cspm-Signature": cspm_sig,
206+
"X-Tokens-Signature": tokens_signature,
137207
"X-Timestamp": tstmp
138208
}
139209

@@ -157,7 +227,7 @@ def update_credentials():
157227

158228
cspm_headers = {"X-API-Key": aqua_api_key, "X-Signature": cspm_sig, "X-Timestamp": tstmp}
159229

160-
cspm_response = http_request(cspm_url + f"/v2/keys/{cspm_key_id}", cspm_headers, "PUT", cspm_body)
230+
cspm_response = cspm_request_with_fallback(cspm_url, f"/v2/keys/{cspm_key_id}", cspm_headers, "PUT", cspm_body, aqua_api_key, aqua_secret, tstmp)
161231

162232
ac_body = json.dumps({
163233
"cloud_account_id": aws_account_id,
@@ -170,8 +240,9 @@ def update_credentials():
170240
})
171241

172242
ac_sig = get_signature(aqua_secret, tstmp, "/v2/internal_apikeys", method="GET")
243+
tokens_signature = get_signature(aqua_secret, tstmp, "/v2/tokens", "POST", '{"validity":1,"allowed_endpoints":["ANY"]}')
173244

174-
ac_headers = {"X-API-Key": aqua_api_key, "X-Authenticate-Api-Key-Signature": ac_sig, "X-Timestamp": tstmp}
245+
ac_headers = {"X-API-Key": aqua_api_key, "X-Authenticate-Api-Key-Signature": ac_sig, "X-Tokens-Signature": tokens_signature, "X-Timestamp": tstmp}
175246

176247
ac_response = http_request(ac_url + f"/discover/update-credentials/{cloud}", ac_headers, "PUT", ac_body)
177248

0 commit comments

Comments
 (0)