def task_finished(client_info): """ Header: {'container_id': abcd...123 'service_name': 'Extract', 'service_version': '4.0.1', 'service_tool_version': ' } Data Block: { "exec_time": 300, "task": <Original Task Dict>, "result": <AL Result Dict>, "freshen": true } """ try: service_name = client_info['service_name'] response = TASKING_CLIENT.task_finished( request.json, client_info['client_id'], service_name, get_metrics_factory(service_name)) if response: return make_api_response(response) return make_api_response("", "No result or error provided by service.", 400) except ValueError as e: # Catch errors when building Task or Result model return make_api_response("", e, 400) except BadRequest: return make_api_response("", "Data received not in JSON format", 400)
def upload_file(client_info): """ Upload a single file. Variables: None Arguments: None Data Block: None Files: Multipart file obj stored in the "file" key. API call example: PUT /api/v1/file/ Result example: {"success": true} """ sha256 = request.headers['sha256'] classification = request.headers['classification'] ttl = int(request.headers['ttl']) is_section_image = request.headers.get('is_section_image', 'false').lower() == 'true' with tempfile.NamedTemporaryFile(mode='bw') as temp_file: # Try reading multipart data from 'files' or a single file post from stream if request.content_type.startswith('multipart'): file = request.files['file'] file.save(temp_file.name) elif request.stream.is_exhausted: if request.stream.limit == len(request.data): temp_file.write(request.data) else: raise ValueError("Cannot find the uploaded file...") else: shutil.copyfileobj(request.stream, temp_file) try: TASKING_CLIENT.upload_file(temp_file.name, classification, ttl, is_section_image, expected_sha256=sha256) except TaskingClientException as e: LOGGER.warning(f"{client_info['client_id']} - {client_info['service_name']}: {str(e)}") return make_api_response(dict(success=False), err=str(e), status_code=400) LOGGER.info(f"{client_info['client_id']} - {client_info['service_name']}: " f"Successfully uploaded file (SHA256: {sha256})") return make_api_response(dict(success=True))
def get_safelisted_tags(**_): """ Get all the safelisted tags in the system Variables: tags => List of tag types (comma seperated) Arguments: None Data Block: None API call example: GET /api/v1/safelist/?tags=network.static.domain,network.dynamic.domain Result example: { "match": { # List of direct matches by tag type "network.static.domain": ["google.ca"], "network.dynamic.domain": ["updates.microsoft.com"] }, "regex": { # List of regular expressions by tag type "network.static.domain": ["*.cyber.gc.ca"], "network.dynamic.domain": ["*.cyber.gc.ca"] } } """ tag_types = request.args.get('tag_types', None) return make_api_response(SAFELIST_CLIENT.get_safelisted_tags(tag_types))
def api_version_list(**kwargs): """ List all available API versions. Variables: None Arguments: None Data Block: None Result example: ["v1"] #List of API versions available """ api_list = [] for rule in current_app.url_map.iter_rules(): if rule.rule.startswith("/api/"): version = rule.rule[5:].split("/", 1)[0] if version not in api_list and version != '': # noinspection PyBroadException try: int(version[1:]) except Exception: continue api_list.append(version) return make_api_response(api_list)
def download_file(sha256, client_info): """ Download a file. Variables: sha256 => A resource locator for the file (sha256) Arguments: None Data Block: None API call example: GET /api/v1/file/123456...654321/ Result example: <THE FILE BINARY> """ with tempfile.NamedTemporaryFile() as temp_file: try: FILESTORE.download(sha256, temp_file.name) f_size = os.path.getsize(temp_file.name) return stream_file_response(open(temp_file.name, 'rb'), sha256, f_size) except FileStoreException: LOGGER.exception(f"[{client_info['client_id']}] {client_info['service_name']} couldn't find file " f"{sha256} requested by service ") return make_api_response({}, "The file was not found in the system.", 404)
def get_task(client_info): """ Header: {'container_id': abcd...123 'service_name': 'Extract', 'service_version': '4.0.1', 'service_tool_version': ' 'timeout': '30'} Result example: {'keep_alive': true} """ service_name = client_info['service_name'] service_version = client_info['service_version'] service_tool_version = client_info['service_tool_version'] client_id = client_info['client_id'] remaining_time = timeout = int(float(request.headers.get('timeout', 30))) metric_factory = get_metrics_factory(service_name) start_time = time.time() status_expiry = start_time + timeout while remaining_time > 0: try: task, retry = TASKING_CLIENT.get_task(client_id, service_name, service_version, service_tool_version, metric_factory, status_expiry=status_expiry, timeout=remaining_time) except ServiceMissingException as e: return make_api_response({}, str(e), 404) if task is not None: return make_api_response(dict(task=task)) elif not retry: return make_api_response(dict(task=False)) # Recalculating how much time we have left before we reach the timeout remaining_time = start_time + timeout - time.time() # We've been processing cache hit for the length of the timeout... bailing out! return make_api_response(dict(task=False))
def register_service(client_info): """ Data Block: < SERVICE MANIFEST > Result example: { 'keep_alive': true, 'new_heuristics': [], 'service_config': < APPLIED SERVICE CONFIG > } """ try: output = TASKING_CLIENT.register_service(request.json, log_prefix=f"{client_info['client_id']} - ") return make_api_response(output) except ValueError as e: return make_api_response("", err=e, status_code=400) except BadRequest: return make_api_response("", "Data received not in JSON format.", 400)
def exists(qhash, **_): """ Check if a file exists in the safelist. Variables: qhash => Hash to check Arguments: None Data Block: None API call example: GET /api/v1/safelist/123456...654321/ Result example: <Safelisting object> """ safelist = SAFELIST_CLIENT.exists(qhash) if safelist: return make_api_response(safelist) return make_api_response(None, "The hash was not found in the safelist.", 404)
def get_safelisted_signatures(**_): """ Get all the signatures that were safelisted in the system. Variables: None Arguments: None Data Block: None API call example: GET /api/v1/safelist/signatures/ Result example: ["McAfee.Eicar", "Avira.Eicar", ...] """ return make_api_response(SAFELIST_CLIENT.get_safelisted_signatures())
def base(*args, **kwargs): # Before anything else, check that the API key is set apikey = request.environ.get('HTTP_X_APIKEY', None) if AUTH_KEY != apikey: client_id = request.headers.get('container_id', 'Unknown Client') header_dump = '; '.join(f'{k}={v}' for k, v in request.headers.items()) wsgi_dump = '; '.join(f'{k}={v}' for k, v in request.environ.items()) LOGGER.warning(f'Client [{client_id}] provided wrong api key [{apikey}] ' f'headers: {header_dump}; wsgi: {wsgi_dump}') return make_api_response("", "Unauthorized access denied", 401) client_info = dict( client_id=request.headers['container_id'], service_name=request.headers['service_name'], service_version=request.headers['service_version'], service_tool_version=request.headers.get('service_tool_version'), ) if config.core.metrics.apm_server.server_url is not None: elasticapm.set_user_context(username=client_info['service_name']) kwargs['client_info'] = client_info return func(*args, **kwargs)