Skip to content

Commit 4a28ad8

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 4a28ad8

File tree

3 files changed

+230
-14
lines changed

3 files changed

+230
-14
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: 86 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,12 @@
66
import http.client
77
import ssl
88

9+
10+
def log(message):
11+
"""Log message to stderr (for Terraform external data source compatibility)"""
12+
print(message, file=sys.stderr)
13+
14+
915
query = json.loads(sys.stdin.read())
1016
ac_url = query.get('autoconnect_url')
1117
cspm_url = query.get('cspm_url')
@@ -49,6 +55,8 @@ def http_request(url, headers, method, body=None):
4955
else:
5056
raise ValueError("Unsupported URL format")
5157

58+
log(f"HTTP request: {method} {url}")
59+
5260
try:
5361
conn = http.client.HTTPSConnection(hostname, context=ssl._create_unverified_context())
5462
conn.request(method, path, body=body, headers=headers)
@@ -58,16 +66,82 @@ def http_request(url, headers, method, body=None):
5866

5967
conn.close()
6068

69+
log(f"HTTP response: {response.status} {response.reason} - {response_data}")
70+
6171
return {
6272
"status": response.status,
6373
"reason": response.reason,
6474
"data": response_data
6575
}
6676
except Exception as e:
67-
print(f"Failed to send HTTP request: {e}")
77+
log(f"Failed to send HTTP request: {e}")
6878
return None
6979

7080

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

@@ -78,15 +152,15 @@ def get_cspm_key_id(aqua_api_key, aqua_secret, cspm_url, role_arn):
78152
"X-Timestamp": tstmp
79153
}
80154

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

83157
if response is None:
84158
raise ValueError(f"HTTP request failed while getting CSPM key ID for {role_arn}")
85159

86160
if response["status"] not in [200, 201]:
87161
raise ValueError(f"Failed to get CSPM key ID: {response['data']}")
88162

89-
json_object = json.loads(response["data"])
163+
json_object = json.loads(response["data"].strip())
90164
for key in json_object['data']:
91165
if key['role_arn'] == role_arn:
92166
return key['id']
@@ -130,10 +204,12 @@ def trigger_discovery():
130204
)
131205

132206
cspm_sig = get_signature(aqua_secret, tstmp, "/v2/keys", "POST", body_cspm)
207+
tokens_signature = get_signature(aqua_secret, tstmp, "/v2/tokens", "POST", '{"validity":1,"allowed_endpoints":["ANY"]}')
133208
headers = {
134209
"X-API-Key": aqua_api_key,
135210
"X-Authenticate-Api-Key-Signature": sig,
136211
"X-Register-New-Cspm-Signature": cspm_sig,
212+
"X-Tokens-Signature": tokens_signature,
137213
"X-Timestamp": tstmp
138214
}
139215

@@ -157,7 +233,7 @@ def update_credentials():
157233

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

160-
cspm_response = http_request(cspm_url + f"/v2/keys/{cspm_key_id}", cspm_headers, "PUT", cspm_body)
236+
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)
161237

162238
ac_body = json.dumps({
163239
"cloud_account_id": aws_account_id,
@@ -170,8 +246,9 @@ def update_credentials():
170246
})
171247

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

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

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

@@ -182,12 +259,13 @@ def update_credentials():
182259

183260

184261
def main():
262+
discovery_response = None
185263
try:
186264
discovery_response = trigger_discovery()
187-
discovery_data = json.loads(discovery_response.get("data", "{}"))
265+
discovery_data = json.loads(discovery_response.get("data", "{}").strip())
188266
if discovery_data.get("state") == "discover_in_progress":
189267
update_credentials_response = update_credentials()
190-
update_credentials_data = json.loads(update_credentials_response.get("data", "{}"))
268+
update_credentials_data = json.loads(update_credentials_response.get("data", "{}").strip())
191269

192270
if update_credentials_response and update_credentials_response.get('status') == 200:
193271
onboarding_status = f"received response: discovery status {discovery_response.get('status', 'Unknown')}, body: {discovery_data}"
@@ -197,7 +275,7 @@ def main():
197275
onboarding_status = f"received response: discovery status {discovery_response.get('status', 'Unknown')}, body: {discovery_data}"
198276

199277
except Exception as e:
200-
onboarding_status = f"received response: status Error, body: {str(e)}"
278+
onboarding_status = f"received response: status Error, body: {str(e)}, raw discovery_response: {discovery_response}"
201279

202280
output = {
203281
"status": onboarding_status

0 commit comments

Comments
 (0)