def test_file_upload(self, mocked_requests_post, mocked_requests_get):
sim_post_content = {}
sim_get_content = {
"results": {
"session": "session_key",
"userId": "1"
}
}
f = "This is file contents"
mocked_requests_post.return_value = self._generateResponse(
sim_post_content, 200)
mocked_requests_get.return_value = self._generateResponse(
sim_get_content, 200)
creds = MockCredsAfterCheck()
file_upload = submit_file(creds, f, "a_new_file.txt")
# Verify call was successful
assert file_upload.status_code == 200
def _mcafee_atd_analyze_file_function(self, event, *args, **kwargs):
"""Function: """
report_file = None
try:
resilient_client = self.rest_client()
inputs = {}
start_time = time.time()
yield StatusMessage("Starting...")
# Get the function parameters:
incident_id = get_incident_id(**kwargs)
artifact_id = kwargs.get("artifact_id") # number
attachment_id = kwargs.get("attachment_id") # number
task_id = kwargs.get("task_id") # number
atd_report_type = self.get_select_param(
kwargs.get("mcafee_atd_report_type")) # select
inputs["macfee_atd_report_type"] = atd_report_type
log.info("incident_id: %s", incident_id)
inputs["incident_id"] = incident_id
if artifact_id is not None:
log.info("artifact_id: %s", artifact_id)
inputs["artifact_id"] = artifact_id
if attachment_id is not None:
log.info("attachment_id: %s", attachment_id)
inputs["attachment_id"] = attachment_id
if task_id is not None:
log.info("task_id: %s", task_id)
inputs["task_id"] = task_id
f_download = _get_file(resilient_client, **kwargs)
f = f_download["file"]
file_name = f_download["file_name"]
yield StatusMessage("File {} downloaded".format(file_name))
response = submit_file(self, f, file_name)
check_status_code(response)
content = response.json()
atd_task_id = content["results"][0]["taskId"]
files_wait = content["filesWait"]
estimated_time = content["estimatedTime"]
yield StatusMessage("File uploaded to ATD with taskId: {}".format(
str(atd_task_id)))
yield StatusMessage("Files waiting on: {}".format(files_wait))
yield StatusMessage(
"Estimated Time: {} minutes".format(estimated_time))
timeout_seconds = self.timeout_mins * 60
start = time.time()
try:
while check_atd_status(self, atd_task_id) is False:
if artifact_id is not None:
yield StatusMessage(
"Analysis on artifact_id {} is still running".
format(artifact_id))
else:
yield StatusMessage(
"Analysis on attachment_id {} is still running".
format(attachment_id))
check_timeout(start, self.polling_interval,
timeout_seconds)
except ValueError:
yield StatusMessage(
"ATD analysis probably failed, please check ATD system.")
raise FunctionError()
yield StatusMessage("Analysis Completed")
if atd_report_type == "pdf" or atd_report_type == "html":
yield StatusMessage(
"Obtaining {} report".format(atd_report_type))
report_file = create_report_file(file_name, atd_report_type)
results = get_atd_report(self, atd_task_id, atd_report_type,
report_file)
if report_file is not None:
resilient_client.post_attachment(
"/incidents/{}/attachments/".format(incident_id),
report_file["report_file"],
filename=report_file["report_file_name"])
yield StatusMessage(
"Report added to incident {} as Attachment".format(
str(incident_id)))
end_time = time.time()
results["Run Time"] = str(end_time - start_time)
results["Inputs"] = inputs
yield StatusMessage("done...")
# Produce a FunctionResult with the results
yield FunctionResult(results)
except Exception as e:
log.info(e)
yield FunctionError("Failed")
finally:
if report_file is not None:
remove_dir(report_file["tmp_dir"])