Skip to content

Commit dd77813

Browse files
[client] report work expectation without impersonation (#10829)
1 parent 67d5d51 commit dd77813

File tree

2 files changed

+23
-17
lines changed

2 files changed

+23
-17
lines changed

pycti/api/opencti_api_client.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -264,13 +264,15 @@ def set_retry_number(self, retry_number):
264264
"" if retry_number is None else str(retry_number)
265265
)
266266

267-
def query(self, query, variables=None):
267+
def query(self, query, variables=None, disable_impersonate=False):
268268
"""submit a query to the OpenCTI GraphQL API
269269
270270
:param query: GraphQL query string
271271
:type query: str
272272
:param variables: GraphQL query variables, defaults to {}
273273
:type variables: dict, optional
274+
:param disable_impersonate: removes impersonate header if set to True, defaults to False
275+
:type disable_impersonate: bool, optional
274276
:return: returns the response json content
275277
:rtype: Any
276278
"""
@@ -295,6 +297,9 @@ def query(self, query, variables=None):
295297
else:
296298
query_var[key] = val
297299

300+
query_headers = self.request_headers.copy()
301+
if disable_impersonate and "opencti-applicant-id" in query_headers:
302+
del query_headers["opencti-applicant-id"]
298303
# If yes, transform variable (file to null) and create multipart query
299304
if len(files_vars) > 0:
300305
multipart_data = {
@@ -361,7 +366,7 @@ def query(self, query, variables=None):
361366
self.api_url,
362367
data=multipart_data,
363368
files=multipart_files,
364-
headers=self.request_headers,
369+
headers=query_headers,
365370
verify=self.ssl_verify,
366371
cert=self.cert,
367372
proxies=self.proxies,
@@ -372,7 +377,7 @@ def query(self, query, variables=None):
372377
r = self.session.post(
373378
self.api_url,
374379
json={"query": query, "variables": variables},
375-
headers=self.request_headers,
380+
headers=query_headers,
376381
verify=self.ssl_verify,
377382
cert=self.cert,
378383
proxies=self.proxies,

pycti/api/opencti_api_work.py

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ def to_received(self, work_id: str, message: str):
2020
}
2121
}
2222
"""
23-
self.api.query(query, {"id": work_id, "message": message})
23+
self.api.query(query, {"id": work_id, "message": message}, True)
2424

2525
def to_processed(self, work_id: str, message: str, in_error: bool = False):
2626
if self.api.bundle_send_to_queue:
@@ -35,7 +35,7 @@ def to_processed(self, work_id: str, message: str, in_error: bool = False):
3535
}
3636
"""
3737
self.api.query(
38-
query, {"id": work_id, "message": message, "inError": in_error}
38+
query, {"id": work_id, "message": message, "inError": in_error}, True
3939
)
4040

4141
def ping(self, work_id: str):
@@ -60,7 +60,7 @@ def report_expectation(self, work_id: str, error):
6060
}
6161
"""
6262
try:
63-
self.api.query(query, {"id": work_id, "error": error})
63+
self.api.query(query, {"id": work_id, "error": error}, True)
6464
except:
6565
self.api.app_logger.error("Cannot report expectation")
6666

@@ -78,7 +78,9 @@ def add_expectations(self, work_id: str, expectations: int):
7878
}
7979
"""
8080
try:
81-
self.api.query(query, {"id": work_id, "expectations": expectations})
81+
self.api.query(
82+
query, {"id": work_id, "expectations": expectations}, True
83+
)
8284
except:
8385
self.api.app_logger.error("Cannot report expectation")
8486

@@ -96,7 +98,9 @@ def add_draft_context(self, work_id: str, draft_context: str):
9698
}
9799
"""
98100
try:
99-
self.api.query(query, {"id": work_id, "draftContext": draft_context})
101+
self.api.query(
102+
query, {"id": work_id, "draftContext": draft_context}, True
103+
)
100104
except:
101105
self.api.app_logger.error("Cannot report draft context")
102106

@@ -111,7 +115,9 @@ def initiate_work(self, connector_id: str, friendly_name: str) -> str:
111115
}
112116
"""
113117
work = self.api.query(
114-
query, {"connectorId": connector_id, "friendlyName": friendly_name}
118+
query,
119+
{"connectorId": connector_id, "friendlyName": friendly_name},
120+
True,
115121
)
116122
return work["data"]["workAdd"]["id"]
117123

@@ -122,10 +128,7 @@ def delete_work(self, work_id: str):
122128
delete
123129
}
124130
}"""
125-
work = self.api.query(
126-
query,
127-
{"workId": work_id},
128-
)
131+
work = self.api.query(query, {"workId": work_id}, True)
129132
return work["data"]
130133

131134
def wait_for_work_to_finish(self, work_id: str):
@@ -179,10 +182,7 @@ def get_work(self, work_id: str) -> Dict:
179182
}
180183
}
181184
"""
182-
result = self.api.query(
183-
query,
184-
{"id": work_id},
185-
)
185+
result = self.api.query(query, {"id": work_id}, True)
186186
return result["data"]["work"]
187187

188188
def get_connector_works(self, connector_id: str) -> List[Dict]:
@@ -243,6 +243,7 @@ def get_connector_works(self, connector_id: str) -> List[Dict]:
243243
"filterGroups": [],
244244
},
245245
},
246+
True,
246247
)
247248
result = result["data"]["works"]["edges"]
248249
return_value = []

0 commit comments

Comments
 (0)