Traps Deprecated

import copy

import urllib3
import jwt

from CommonServerPython import *

# Disable insecure warnings
urllib3.disable_warnings()

PARAMS = demisto.params()
# Remove trailing slash to prevent wrong URL path to service
SERVER = PARAMS['url'][:-1] \
    if (PARAMS['url'] and PARAMS['url'].endswith('/')) else PARAMS['url']
# Should we use SSL
USE_SSL = not PARAMS.get('insecure', False)
# Service base URL
BASE_URL = SERVER + '/xapi/v1/'
APPLICATION_ID = PARAMS.get('application_id')
PRIVATE_KEY = PARAMS.get('private_key')
# Headers to be sent in requests
REQUEST_HEADERS = {
    'Content-Type': 'application/json'
}

OUTPUTS = {
    'get_endpoint_by_id': {
        'ID': 'guid',
        'Name': 'name',
        'Domain': 'domain',
        'Platform': 'platform',
        'ScanStatus': 'scanStatus',
        'Status': 'status',
        'IP': 'ip',
        'ComputerSid': 'computerSid',
        'IsCompromised': 'compromised',
        'OsVersion': 'osVersion',
        'OsProductType': 'osProductType',
        'OsProductName': 'osProductName',
        'Is64': 'is64',
        'LastSeen': 'lastSeen',
        'LastUser': 'lastUser'
    },
    'endpoint_files_retrieve': {
        'OperationID': 'operationId'
    },
    'endpoint_isolate': {
        'OperationID': 'operationId'
    },
    'endpoint_scan': {
        'OperationID': 'operationId'
    },
    'event_bulk_update_status': {
        'EventID': 'eventGuid'
    },
    'hashes_blacklist_status': {
        'SHA256': 'hash',
        'BlacklistStatus': 'status'
    },
    'event_quarantine_result': {
        'SHA256': 'fileHash',
        'FilePath': 'filePath'
    },
    'endpoint_scan_result': {
        'FileScanned': 'filesScanned',
        'FilesFailed': 'filesFailed',
        'MalwareFound': 'malwareFound'
    }
}


def create_headers(with_auth):
    """Create headers for the http_request

    Args:
        with_auth: True

    Returns:
        Headers.
    """
    headers = copy.deepcopy(REQUEST_HEADERS)
    if with_auth:
        token = generate_auth_token().decode('utf-8')
        headers['Authorization'] = f'Bearer {token}'
    return headers


def extract_and_validate_http_response(response, operation_err_message, test=False):
    """

    Args:
        response: raw response
        operation_err_message: error message to present in case of error
        test: boolean value, true if test

    Returns:
        Error if response is faulty.
    """
    try:
        response.raise_for_status()
        return response.json() if not test else response.content
    except requests.exceptions.HTTPError:
        try:
            err_message = response.json().get('message')
        except Exception:
            try:
                err_obj = json.loads(xml2json(response.text))
                err_message = demisto.get(err_obj, 'Error.Message')
            except Exception:
                err_message = 'Could not parse error'
        return_error(f'{operation_err_message}: \n{err_message}')


def http_request(method, url_suffix, plain_url=False, params=None, data=None, operation_err=None, parse_response=True,
                 with_auth=True):
    """Generic http call to Traps

    Args:
        method: request method.
        url_suffix: URL suffix.
        plain_url: full URL.
        params: request params.
        data: request data.
        operation_err: operation error to log the user.
        parse_response: boolean value, if parsing the response is needed.
        with_auth: boolean value, do we need to authenticate the request.

    Returns:
        Result from the API.
    """
    try:
        result = requests.request(
            method,
            BASE_URL + url_suffix if not plain_url else url_suffix,
            verify=USE_SSL,
            params=params,
            data=json.dumps(data) if data else data,
            headers=create_headers(with_auth),
        )
    except requests.exceptions.ConnectionError:
        return_error('Error connecting to Traps server. Please check your connection and you server address')
    if parse_response:
        result = extract_and_validate_http_response(result, operation_err, plain_url)
    return result


def health_check():
    """Performs basic health check on the server.

    Returns:
        Error if not ok.
    """
    path = f'{SERVER}/xapi/health-check'
    server_status = http_request('GET', path, plain_url=True).decode('utf-8')
    if server_status == '"Ok"':
        return
    raise Exception(f'Server health-check failed. Status returned was: {server_status}')


def generate_auth_token():
    """Generate a token using jwt.

    Returns:
        token.
    """
    key = PRIVATE_KEY
    data = {'appId': APPLICATION_ID}
    token = jwt.encode(data, key, algorithm='RS256')
    return token


def parse_data_from_response(resp_obj, operation_name=None):
    """Response raw data.

    Args:
        resp_obj: raw_data.
        operation_name: operation name.

    Returns:
        parsed data.
    """
    new_data_obj = {}  # type: dict
    outputs_obj = OUTPUTS[operation_name]
    for key, val in outputs_obj.items():
        if val in resp_obj:
            new_data_obj[key] = resp_obj.get(val)

    return new_data_obj


def get_endpoint_by_id(endpoint_id):
    """Get endpoint data by sending a GET request.

    Args:
        endpoint_id: endpoint ID.

    Returns:
        endpoint data.
    """
    path = f'agents/{endpoint_id}'
    endpoint_data = http_request('GET', path, operation_err=f'Get endpoint {endpoint_id} failed')
    return parse_data_from_response(endpoint_data, 'get_endpoint_by_id'), endpoint_data


def endpoint_files_retrieve(endpoint_id, file_name, event_id):
    """Retrieve a file from the endpoint by sending a POST request.

    Args:
        endpoint_id: endpoint ID.
        file_name: File name.
        event_id: Event ID.

    Returns:
        Operation data.
    """
    path = f'agents/{endpoint_id}/files-retrieve'
    data = {
        'incidentId': event_id,
        'files': [
            {
                "path": file_name
            }
        ]
    }
    resp = http_request('POST', path, data=data,
                        operation_err=f'Files retrieve command on endpoint {endpoint_id} failed')
    operation_obj = parse_data_from_response(resp, 'endpoint_files_retrieve')
    operation_obj.update({
        'EndpointID': endpoint_id,
        'Type': 'files-retrieve'
    })
    return operation_obj


def endpoint_scan(endpoint_id):
    """Initiate a scan on an endpoint by sending a POST request.

    Args:
        endpoint_id: endpoint ID.

    Returns:
        Operation data.
    """
    path = f'agents/{endpoint_id}/scan'
    response = http_request('POST', path, operation_err=f'Scanning endpoint: {endpoint_id} failed')
    operation_obj = parse_data_from_response(response, 'endpoint_scan')
    operation_obj.update({
        'EndpointID': endpoint_id,
        'Type': 'endpoint-scan'
    })
    return operation_obj, response


def endpoint_scan_result(operation_id):
    """Initiate the SAM operation for retrieving the scan result.

    Args:
        operation_id: operation ID.

    Returns:
        scan data.
    """
    status, additional_data = sam_operation(operation_id, 'Could not get scan results')
    scan_data = parse_data_from_response(additional_data.get('scanData'),
                                         'endpoint_scan_result') if additional_data else {}
    scan_data['Status'] = status
    scan_data['OperationID'] = operation_id
    return scan_data


def update_event_status(event_ids, status):
    """Update event or events status by sending a POST request.

    Args:
        event_ids: event IDs.
        status: status.

    Returns:
        API response.
    """
    path = 'events/status'
    data = {
        "guids": event_ids,
        "status": status
    }
    response = http_request('PATCH', path, data=data, operation_err=f'Update events {event_ids} status failed')
    return response


def update_event_comment(event_id, comment):
    """Update event comment by sending a POST request.

    Args:
        event_id: event ID.
        comment: comment.
    """
    path = f'events/{event_id}/comment'
    data = {
        "comment": comment
    }
    http_request('POST', path, data=data, operation_err=f'Update event: {event_id} comment failed')


def event_update(event_id, status, comment):
    """Initiate update of an event.

    Args:
        event_id: event ID.
        status: status.
        comment: comment.

    Returns:
        Error if not successful.
    """
    if not status and not comment:
        raise Exception('Please add a status or a comment. Neither was given.')
    if status:
        response = update_event_status([event_id], status)
        if response.get('failed'):
            raise Exception(f'Update status for event: {event_id} has failed.')
    if comment:
        update_event_comment(event_id, comment)


def event_bulk_update_status(event_ids, status):
    """Initiate update of events statuses.

    Args:
        event_ids: event IDs.
        status: status.

    Returns:
        Update statuses.
    """
    ids_obj = update_event_status(event_ids, status)
    results = {
        'UpdateSuccess': list(
            map(lambda id_obj: parse_data_from_response(id_obj, 'event_bulk_update_status'), ids_obj.get('succeeded'))),
        'UpdateFail': list(
            map(lambda id_obj: parse_data_from_response(id_obj, 'event_bulk_update_status'), ids_obj.get('failed'))),
        'UpdateIgnored': list(
            map(lambda id_obj: parse_data_from_response(id_obj, 'event_bulk_update_status'), ids_obj.get('ignored')))
    }
    return results


def hash_blacklist(hash_id):
    """Add hash to a blacklist by sending a POST request.

    Args:
        hash_id: Hash.

    Returns:
        status.
    """
    path = f'hashes/{hash_id}/blacklist'
    result = http_request('POST', path, operation_err=f'Failed to blacklist {hash_id}')
    return result.get('status')


def remove_hash_from_blacklist(hash_id):
    """Remove hash to a blacklist by sending a POST request.

    Args:
        hash_id: Hash.

    Returns:
        status.
    """
    path = f'hashes/{hash_id}/blacklist-remove'
    result = http_request('POST', path, operation_err=f'Failed to remove {hash_id} from blacklist')
    return result.get('status')


def hashes_blacklist_status(hash_ids):
    """Get hashes blacklisting status by sending a POST request.

    Args:
        hash_ids: Hashes.

    Returns:
        Hashes and blacklisting data.
    """
    path = 'hashes/blacklist-status'
    data = {
        'hashes': hash_ids
    }
    ids_obj = http_request('POST', path, data=data, operation_err='Failed to get hashes status')
    result = list(map(lambda id_obj: parse_data_from_response(id_obj, 'hashes_blacklist_status'), ids_obj))
    return result


def event_quarantine(event_id):
    """Quarantine an event by sending a POST request.

    Args:
        event_id: event ID.

    Returns:
        Data regarding the event and the quarantine operation.
    """
    path = f'events/{event_id}/quarantine'
    resp = http_request('POST', path, operation_err=f'Quarantine event {event_id} failed')
    message_ids = resp.get('operationId').get('samMessageIds')
    operations = []
    for op_id in message_ids:
        operations.append({
            'EventID': event_id,
            'Type': 'event-quarantine',
            'OperationID': op_id
        })
    return operations


def endpoint_isolate(endpoint_id):
    """Isolate an endpoint by sending a POST request.

    Args:
        endpoint_id: endpoint ID.

    Returns:
        Data regarding the endpoint and the isolation operation.
    """
    path = f'agents/{endpoint_id}/isolate'
    resp = http_request('POST', path, operation_err=f'Isolation of endpoint: {endpoint_id} failed')
    operation_obj = parse_data_from_response(resp, 'endpoint_isolate')
    operation_obj.update({
        'EndpointID': endpoint_id,
        'Type': 'endpoint-isolate'
    })
    return operation_obj


def sam_operation(operation_id, operation_err):
    """
    This functions invokes an API call to the sam operation endpoint on Traps server to get the operation status and/or
    results.
    :param operation_id: the operation on which to get the status/results
    :param operation_err: The error to return in case of a failure (changes according to the command fired.)
    :return:
        status: the status of the operation.
        additional_data: additional data regarding the operation (like scan results)
    """
    path = f'sam/operations/{operation_id}'
    result = http_request('GET', path, operation_err=operation_err)
    summary_data = result.get('summaryData')
    if summary_data and (summary_data.get('incompatible') or summary_data.get('samExists')):
        if operation_err == 'Could not get scan results':  # Get scan result
            requested_scans = int(summary_data.get('requested', 0))
            incompatible_scans = int(summary_data.get('incompatible', 0))
            sam_exists_scans = int(summary_data.get('samExists', 0))
            if requested_scans <= incompatible_scans + sam_exists_scans:
                raise Exception(f'{operation_err}.\nRequested scans number: {requested_scans}.\n'
                                f'Incompatible scans number: {requested_scans}.\n'
                                f'Sam exists scans number: {sam_exists_scans}.')
        raise Exception(f'{operation_err}')
    if result.get('summaryData').get('samExists'):
        return 'ignored', None
    for status_obj in result.get('statuses'):
        if status_obj.get('count') > 0:
            return status_obj.get('status'), result.get('additionalData')
    raise Exception(f'{operation_err}: Could not retrieve status')


def endpoint_isolate_status(operation_id):
    """Initiate the SAM operation for endpoint isolation status.

    Args:
        operation_id: operation ID.

    Returns:
        endpoint status, operation ID
    """
    status, _ = sam_operation(operation_id, 'Could not get endpoint isolate status')
    return {'Status': status, 'OperationID': operation_id}


def event_quarantine_result(operation_id):
    """Initiate the SAM operation for quarantine result.

    Args:
        operation_id: operation ID.

    Returns:
        quarantine data.
    """
    status, additional_data = sam_operation(operation_id, 'Could not get event quarantine status')
    quarantine_data = parse_data_from_response(additional_data.get('quarantineData'),
                                               'event_quarantine_result') if additional_data else {}
    quarantine_data['Status'] = status
    quarantine_data['OperationID'] = operation_id
    return quarantine_data


def endpoint_files_retrieve_result(operation_id):
    """Initiate the SAM operation for file retrieving.

    Args:
        operation_id: operation ID.

    Returns:
        file as an entry, status, operation_id.
    """
    status, additional_data = sam_operation(operation_id, 'Failed to get file retrieve results')
    if status == 'finished':
        file_info = additional_data.get('uploadData')
        file_name = file_info.get('fileName')
        url = file_info.get('downloadUrl')
        data = http_request('GET', url, plain_url=True, operation_err='Unable to download file.', with_auth=False)
        demisto.results(fileResult(filename=file_name, data=data))
    return {'Status': status, 'OperationID': operation_id}


def test_module_command():
    """Performs basic GET request to check if the API is reachable and authentication is successful.

    Returns:
        ok if successful.
    """
    health_check()
    result = http_request('GET', 'agents/1', parse_response=False)
    if result.status_code == 403:
        raise Exception('Error connecting to server. Check your Application ID and Private key')
    demisto.results('ok')


def get_endpoint_by_id_command():
    """Get endpoint data.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    endpoint_id = args.get('endpoint_id')
    endpoint_data, raw_data = get_endpoint_by_id(endpoint_id)
    human_readable = tableToMarkdown(f'Endpoint {endpoint_id} data:', endpoint_data, headerTransform=pascalToSpace)
    context = {'Traps.Endpoint(val.ID == obj.ID)': createContext(endpoint_data)}
    return_outputs(human_readable, context, raw_response=raw_data)


def endpoint_files_retrieve_command():
    """Initiate retrieving of files from an endpoint.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    endpoint_id = args.get('endpoint_id')
    file_name = args.get('file_name')
    event_id = args.get('event_id')
    operation_obj = endpoint_files_retrieve(endpoint_id, file_name, event_id)
    human_readable = tableToMarkdown(f'Files retrieve command on endpoint: {endpoint_id} received', operation_obj,
                                     headerTransform=pascalToSpace)
    context = {'Traps.FileRetrieve(val.OperationID == obj.OperationID)': operation_obj}
    return_outputs(human_readable, context, operation_obj)


def endpoint_files_retrieve_result_command():
    """Retrieve files from an endpoint result.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    operation_id = args.get('operation_id')
    status_obj = endpoint_files_retrieve_result(operation_id)
    human_readable = f'### File retrieval status is: {status_obj.get("Status")}'
    context = {'Traps.FileRetrieveResult(val.OperationID == obj.OperationID)': status_obj}
    return_outputs(human_readable, context)


def endpoint_scan_command():
    """Initiate scan on an endpoint.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    endpoint_id = args.get('endpoint_id')

    # check that running a scan is possible
    _, raw_data = get_endpoint_by_id(endpoint_id)
    scan_status = raw_data.get('scanStatus')
    if scan_status and scan_status in ['pending', 'in_progress']:
        raise Exception(f'Could not initiate a scan on the endpoint {endpoint_id}'
                        f' because endpoint scan status is {scan_status}.')

    operation_obj, raw_data = endpoint_scan(endpoint_id)
    human_readable = tableToMarkdown(f'Scan command on endpoint: {endpoint_id} received', operation_obj,
                                     headerTransform=pascalToSpace)
    context = {'Traps.Scan(val.OperationID == obj.OperationID)': operation_obj}
    return_outputs(human_readable, context, raw_data)


def endpoint_scan_result_command():
    """Retrieve endpoint scan results.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    operation_id = args.get('operation_id')
    status_obj = endpoint_scan_result(operation_id)
    context = {'Traps.ScanResult(val.OperationID == obj.OperationID)': status_obj}
    human_readable = tableToMarkdown(f'Status of scan operation: {operation_id}', status_obj,
                                     headerTransform=pascalToSpace)
    return_outputs(human_readable, context, status_obj)


def event_update_command():
    """Update an event.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    event_id = args.get('event_id')
    status = args.get('status')
    comment = args.get('comment')
    event_update(event_id, status, comment)
    human_readable = f'### Event: {event_id} was updated'
    human_readable += f'\n##### New status: {status}' if status else ''
    human_readable += f'\n##### New comment: {comment}' if comment else ''
    return_outputs(human_readable, None, {})


def event_bulk_update_status_command():
    """Update events.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    event_ids = argToList(args.get('event_ids'))
    status = args.get('status')
    results = event_bulk_update_status(event_ids, status)
    human_readable = tableToMarkdown('Successfully updated', results.get('UpdateSuccess'),
                                     headerTransform=pascalToSpace)
    human_readable += tableToMarkdown('Failed to update', results.get('UpdateFail'), headerTransform=pascalToSpace)
    human_readable += tableToMarkdown('Ignored', results.get('UpdateIgnored'), headerTransform=pascalToSpace)
    return_outputs(human_readable, {}, {})


def event_quarantine_command():
    """Quarantine an event.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    event_id = args.get('event_id')
    operations = event_quarantine(event_id)
    human_readable = tableToMarkdown(f'Quarantine command on event: {event_id} received', operations,
                                     headerTransform=pascalToSpace)
    context = {'Traps.Quarantine(val.OperationID == obj.OperationID)': operations}
    return_outputs(human_readable, context, operations)


def event_quarantine_result_command():
    """Check quarantine event status.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    operation_id = args.get('operation_id')
    status_obj = event_quarantine_result(operation_id)
    context = {'Traps.QuarantineResult(val.OperationID == obj.OperationID)': status_obj}
    human_readable = tableToMarkdown(f'Status of quarantine operation: {operation_id}',
                                     status_obj, headerTransform=pascalToSpace)
    return_outputs(human_readable, context, status_obj)


def hash_blacklist_command():
    """Add a hash to a blacklist.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    hash_id = args.get('hash_id')
    status = hash_blacklist(hash_id)
    context = {}  # type: dict
    if status == 'success':
        human_readable = f'#### Successfully blacklisted: {hash_id}'
        status_obj = {
            'SHA256': hash_id,
            'BlacklistStatus': 'blacklisted'
        }
        context = {'Traps.File(val.SHA256 == obj.SHA256)': status_obj}

    elif status == 'ignore':
        human_readable = f'#### Hash: {hash_id} already appears in blacklist'
    else:
        human_readable = f'#### Failed to blacklist: {hash_id}'
    return_outputs(human_readable, context, status)


def hash_blacklist_remove_command():
    """Remove a hash from blacklist.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    hash_id = args.get('hash_id')
    status = remove_hash_from_blacklist(hash_id)
    context = {}  # type: dict
    if status == 'success':
        human_readable = f'#### Successfully removed {hash_id} from blacklist'
        status_obj = {
            'SHA256': hash_id,
            'BlacklistStatus': 'none'
        }
        context = {'Traps.File(val.SHA256 == obj.SHA256)': status_obj}
    else:
        human_readable = f'#### Failed to remove {hash_id} from blacklist:'

    return_outputs(human_readable, context, status)


def hashes_blacklist_status_command():
    """Check hash blacklist status.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    hash_ids = args.get('hash_ids').split(',')
    ids_obj = hashes_blacklist_status(hash_ids)
    human_readable = tableToMarkdown('Hashes status:', ids_obj, headerTransform=pascalToSpace)
    context = {'Traps.File(val.SHA256 == obj.SHA256)': ids_obj}
    return_outputs(human_readable, context, ids_obj)


def endpoint_isolate_command():
    """Isolate an endpoint.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    endpoint_id = args.get('endpoint_id')
    operation_obj = endpoint_isolate(endpoint_id)
    human_readable = tableToMarkdown(f'Isolate command on endpoint {endpoint_id} received', operation_obj,
                                     headerTransform=pascalToSpace)
    context = {'Traps.Isolate(val.OperationID == obj.OperationID)': operation_obj}
    return_outputs(human_readable, context, operation_obj)


def endpoint_isolate_status_command():
    """Check endpoint isolation status.

    Returns:
        Demisto Outputs.
    """
    args = demisto.args()
    operation_id = args.get('operation_id')
    isolate_status = endpoint_isolate_status(operation_id)
    human_readable = f'### Isolate status is: {isolate_status.get("Status")}'
    context = {'Traps.IsolateResult(val.OperationID == obj.OperationID)': isolate_status}
    return_outputs(human_readable, context, isolate_status)


def main():
    """
    Initiate integration command
    """
    # Remove proxy if not set to true in params
    handle_proxy()
    command = demisto.command()
    LOG(f'Command being called is {command}.')
    try:
        if command == 'test-module':
            test_module_command()
        elif command == 'traps-get-endpoint-by-id':
            get_endpoint_by_id_command()
        elif command == 'traps-endpoint-files-retrieve':
            endpoint_files_retrieve_command()
        elif command == 'traps-endpoint-files-retrieve-result':
            endpoint_files_retrieve_result_command()
        elif command == 'traps-endpoint-scan':
            endpoint_scan_command()
        elif command == 'traps-endpoint-scan-result':
            endpoint_scan_result_command()
        elif command == 'traps-event-update':
            event_update_command()
        elif command == 'traps-event-bulk-update-status':
            event_bulk_update_status_command()
        elif command == 'traps-hash-blacklist':
            hash_blacklist_command()
        elif command == 'traps-hash-blacklist-remove':
            hash_blacklist_remove_command()
        elif command == 'traps-hashes-blacklist-status':
            hashes_blacklist_status_command()
        elif command == 'traps-event-quarantine':
            event_quarantine_command()
        elif command == 'traps-event-quarantine-result':
            event_quarantine_result_command()
        elif command == 'traps-endpoint-isolate':
            endpoint_isolate_command()
        elif command == 'traps-endpoint-isolate-status':
            endpoint_isolate_status_command()
        else:
            raise NotImplementedError(f'Command {command} was not implemented.')
    except Exception as err:
        return_error(err)


if __name__ in ["__builtin__", "builtins", "__main__"]:
    main()