"""
This playbook investigates and remediates malware infections on the endpoint.
"""

import phantom.rules as phantom
import json
from datetime import datetime, timedelta
##############################
# Start - Global Code Block

"""Malicous file detected on endpoint"""

# End - Global Code block
##############################

def on_start(container):
    phantom.debug('on_start() called')
    
    # call 'file_reputation_1' block
    file_reputation_1(container=container)

    return

def filter_3(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('filter_3() called')

    # collect filtered artifact ids for 'if' condition 1
    matched_artifacts_1, matched_results_1 = phantom.condition(
        container=container,
        action_results=results,
        conditions=[
            ["file_reputation_1:action_result.summary.positives", ">", 5],
            ["file_reputation_1:action_result.summary.positives", "<=", 10],
        ],
        logical_operator='and',
        name="filter_3:condition_1")

    # call connected blocks if filtered artifacts or results
    if matched_artifacts_1 or matched_results_1:
        create_ticket_3(action=action, success=success, container=container, results=results, handle=handle, custom_function=custom_function, filtered_artifacts=matched_artifacts_1, filtered_results=matched_results_1)

    return

def shutdown_system_1(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('shutdown_system_1() called')
        
    #phantom.debug('Action: {0} {1}'.format(action['name'], ('SUCCEEDED' if success else 'FAILED')))
    
    # collect data for 'shutdown_system_1' call
    results_data_1 = phantom.collect2(container=container, datapath=['hunt_file_2:action_result.data.*.process.results.*.hostname', 'hunt_file_2:action_result.parameter.context.artifact_id'], action_results=results)

    parameters = []
    
    # build parameters list for 'shutdown_system_1' call
    for results_item_1 in results_data_1:
        parameters.append({
            'ph': "",
            'message': "",
            'wait_time': "",
            'ip_hostname': results_item_1[0],
            # context (artifact id) is added to associate results with the artifact
            'context': {'artifact_id': results_item_1[1]},
        })

    phantom.act(action="shutdown system", parameters=parameters, assets=['domainctrl1'], callback=join_filter_2, name="shutdown_system_1", parent_action=action)

    return

def create_ticket_2(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    
    #phantom.debug('Action: {0} {1}'.format(action['name'], ('SUCCEEDED' if success else 'FAILED')))
    
    # collect data for 'create_ticket_2' call

    disabled_users = set(phantom.collect2(datapath='disable_user_1:action_result.parameter.username'))
    blocked_hashes = set(phantom.collect2(datapath='block_hash_3:action_result.parameter.hash'))
    loggedoff_users = set(phantom.collect2(datapath='logoff_user_1:action_result.parameter.username'))
    shutdown_systems = set(phantom.collect2(datapath='shutdown_system_1:action_result.parameter.ip_hostname'))
    file_reputation = phantom.collect2(datapath=['file_reputation_1:filtered-action_result.parameter.hash',                                                         
                                                 'file_reputation_1:filtered-action_result.summary.positives'])
    detected_users = set(phantom.collect2(datapath='hunt_file_2:action_result.data.*.process.results.*.username'))
    detected_systems = set(phantom.collect2(datapath='hunt_file_2:action_result.data.*.process.results.*.hostname'))
    
    title = "Virus Detected on {0} devices".format(len(detected_systems))
    
    description =  "Hashes sumbitted with detections:\n{0}\n\n".format(", ".join(["{0} ({1})".format(*fr) for fr in file_reputation]))
    description += "File was found on {0} devices:\n{1}\n\n".format(len(detected_systems), ', '.join(detected_systems))
    description += "This impacts at least {0} users:\n{1}\n\n".format(len(detected_users), ', '.join(detected_users))
    if len(blocked_hashes):
        description += "{0} hashes were submitted for blocking:\n{1}\n\n".format(len(blocked_hashes), ", ".join(blocked_hashes))
    if len(loggedoff_users):
        description += "{0} users were forced to logoff:\n{1}\n\n".format(len(loggedoff_users), ", ".join(loggedoff_users))
    if len(disabled_users):
        description += "{0} user accounts were disabled:\n{1}\n\n".format(len(disabled_users), ", ".join(disabled_users))
    if len(shutdown_systems):
        description += "{0} systems were shutdown:\n{1}\n\n".format(len(shutdown_systems), ", ".join(shutdown_systems))

    parameters = []
    
    # build parameters list for 'create_ticket_2' call
    parameters.append({
        'short_description': title,
        'description': description,
        'fields': "",
    })

    if parameters:
        phantom.act("create ticket", parameters=parameters, assets=['servicenow'], name="create_ticket_2", parent_action=action)    
    else:
        phantom.error("'create_ticket_2' will not be executed due to lack of parameters")
    
    return

def filter_2(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('filter_2() called')

    # collect filtered artifact ids for 'if' condition 1
    matched_artifacts_1, matched_results_1 = phantom.condition(
        container=container,
        action_results=results,
        conditions=[
            ["file_reputation_1:action_result.summary.positives", ">", 10],
        ],
        name="filter_2:condition_1")

    # call connected blocks if filtered artifacts or results
    if matched_artifacts_1 or matched_results_1:
        create_ticket_2(action=action, success=success, container=container, results=results, handle=handle, custom_function=custom_function, filtered_artifacts=matched_artifacts_1, filtered_results=matched_results_1)

    return

def join_filter_2(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None):
    phantom.debug('join_filter_2() called')

    # check if all connected incoming playbooks, actions, or custom functions are done i.e. have succeeded or failed
    if phantom.completed(action_names=['logoff_user_1', 'shutdown_system_1', 'disable_user_1', 'block_hash_3']):
        
        # call connected block "filter_2"
        filter_2(container=container, handle=handle)
    
    return

def logoff_user_1(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('logoff_user_1() called')
        
    #phantom.debug('Action: {0} {1}'.format(action['name'], ('SUCCEEDED' if success else 'FAILED')))
    
    # collect data for 'logoff_user_1' call
    results_data_1 = phantom.collect2(container=container, datapath=['hunt_file_2:action_result.data.*.process.results.*.hostname', 'hunt_file_2:action_result.parameter.context.artifact_id'], action_results=results)

    parameters = []
    
    # build parameters list for 'logoff_user_1' call
    for results_item_1 in results_data_1:
        parameters.append({
            'username': "",
            'ip_hostname': results_item_1[0],
            # context (artifact id) is added to associate results with the artifact
            'context': {'artifact_id': results_item_1[1]},
        })

    phantom.act(action="logoff user", parameters=parameters, assets=['domainctrl1'], callback=join_filter_2, name="logoff_user_1", parent_action=action)

    return

def filter_1(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('filter_1() called')

    # collect filtered artifact ids for 'if' condition 1
    matched_artifacts_1, matched_results_1 = phantom.condition(
        container=container,
        action_results=results,
        conditions=[
            ["file_reputation_1:action_result.summary.positives", ">", 5],
            ["file_reputation_1:action_result.summary.positives", "<=", 10],
        ],
        logical_operator='and',
        name="filter_1:condition_1")

    # call connected blocks if filtered artifacts or results
    if matched_artifacts_1 or matched_results_1:
        hunt_file_1(action=action, success=success, container=container, results=results, handle=handle, custom_function=custom_function, filtered_artifacts=matched_artifacts_1, filtered_results=matched_results_1)
        get_file_2(action=action, success=success, container=container, results=results, handle=handle, custom_function=custom_function, filtered_artifacts=matched_artifacts_1, filtered_results=matched_results_1)

    # collect filtered artifact ids for 'if' condition 2
    matched_artifacts_2, matched_results_2 = phantom.condition(
        container=container,
        action_results=results,
        conditions=[
            ["file_reputation_1:action_result.summary.positives", ">", 10],
        ],
        name="filter_1:condition_2")

    # call connected blocks if filtered artifacts or results
    if matched_artifacts_2 or matched_results_2:
        hunt_file_2(action=action, success=success, container=container, results=results, handle=handle, custom_function=custom_function, filtered_artifacts=matched_artifacts_2, filtered_results=matched_results_2)
        get_file_3(action=action, success=success, container=container, results=results, handle=handle, custom_function=custom_function, filtered_artifacts=matched_artifacts_2, filtered_results=matched_results_2)

    return

def hunt_file_2(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('hunt_file_2() called')
        
    #phantom.debug('Action: {0} {1}'.format(action['name'], ('SUCCEEDED' if success else 'FAILED')))
    
    # collect data for 'hunt_file_2' call
    filtered_results_data_1 = phantom.collect2(container=container, datapath=["filtered-data:filter_1:condition_2:file_reputation_1:action_result.parameter.hash", "filtered-data:filter_1:condition_2:file_reputation_1:action_result.parameter.context.artifact_id"])

    parameters = []
    
    # build parameters list for 'hunt_file_2' call
    for filtered_results_item_1 in filtered_results_data_1:
        if filtered_results_item_1[0]:
            parameters.append({
                'hash': filtered_results_item_1[0],
                'type': "",
                'range': "",
                # context (artifact id) is added to associate results with the artifact
                'context': {'artifact_id': filtered_results_item_1[1]},
            })

    phantom.act(action="hunt file", parameters=parameters, assets=['carbonblack'], callback=hunt_file_2_callback, name="hunt_file_2")

    return

def hunt_file_2_callback(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None):
    phantom.debug('hunt_file_2_callback() called')
    
    disable_user_1(action=action, success=success, container=container, results=results, handle=handle, custom_function=custom_function)
    logoff_user_1(action=action, success=success, container=container, results=results, handle=handle, custom_function=custom_function)
    shutdown_system_1(action=action, success=success, container=container, results=results, handle=handle, custom_function=custom_function)
    block_hash_3(action=action, success=success, container=container, results=results, handle=handle, custom_function=custom_function)

    return

def get_file_3(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('get_file_3() called')
        
    #phantom.debug('Action: {0} {1}'.format(action['name'], ('SUCCEEDED' if success else 'FAILED')))
    
    # collect data for 'get_file_3' call
    filtered_results_data_1 = phantom.collect2(container=container, datapath=["filtered-data:filter_1:condition_2:file_reputation_1:action_result.parameter.hash", "filtered-data:filter_1:condition_2:file_reputation_1:action_result.parameter.context.artifact_id"])

    parameters = []
    
    # build parameters list for 'get_file_3' call
    for filtered_results_item_1 in filtered_results_data_1:
        parameters.append({
            'hash': filtered_results_item_1[0],
            'ph_0': "",
            'offset': "",
            'get_count': "",
            'sensor_id': "",
            'file_source': "",
            # context (artifact id) is added to associate results with the artifact
            'context': {'artifact_id': filtered_results_item_1[1]},
        })

    phantom.act(action="get file", parameters=parameters, assets=['carbonblack'], name="get_file_3")

    return

def get_file_2(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('get_file_2() called')
        
    #phantom.debug('Action: {0} {1}'.format(action['name'], ('SUCCEEDED' if success else 'FAILED')))
    
    # collect data for 'get_file_2' call
    filtered_results_data_1 = phantom.collect2(container=container, datapath=["filtered-data:filter_1:condition_1:file_reputation_1:action_result.parameter.hash", "filtered-data:filter_1:condition_1:file_reputation_1:action_result.parameter.context.artifact_id"])

    parameters = []
    
    # build parameters list for 'get_file_2' call
    for filtered_results_item_1 in filtered_results_data_1:
        parameters.append({
            'hash': filtered_results_item_1[0],
            'ph_0': "",
            'offset': "",
            'get_count': "",
            'sensor_id': "",
            'file_source': "",
            # context (artifact id) is added to associate results with the artifact
            'context': {'artifact_id': filtered_results_item_1[1]},
        })

    phantom.act(action="get file", parameters=parameters, assets=['carbonblack'], name="get_file_2")

    return

def create_ticket_3(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    
    #phantom.debug('Action: {0} {1}'.format(action['name'], ('SUCCEEDED' if success else 'FAILED')))
    
    # collect data for 'create_ticket_2' call

    file_reputation = phantom.collect2(datapath=['file_reputation_1:filtered-action_result.parameter.hash',                                                         
                                                 'file_reputation_1:filtered-action_result.summary.positives'])
    blocked_hashes = set(phantom.collect2(datapath='block_hash_2:action_result.parameter.hash'))
    detected_users = set(phantom.collect2(datapath='hunt_file_1:action_result.data.*.process.results.*.username'))
    detected_systems = set(phantom.collect2(datapath='hunt_file_1:action_result.data.*.process.results.*.hostname'))
    
    title = "Virus Detected on {0} devices".format(len(detected_systems))
    
    description =  "Hashes sumbitted with detections:\n{0}\n\n".format(", ".join(["{0} ({1})".format(*fr) for fr in file_reputation]))
    description += "File was found on {0} devices:\n{1}\n\n".format(len(detected_systems), ', '.join(detected_systems))
    description += "This impacts at least {0} users:\n{1}\n\n".format(len(detected_users), ', '.join(detected_users))
    if len(blocked_hashes):
        description += "{0} hashes were submitted for blocking:\n{1}\n\n".format(len(blocked_hashes), ", ".join(blocked_hashes))

    parameters = []
    
    # build parameters list for 'create_ticket_2' call
    parameters.append({
        'short_description': title,
        'description': description,
        'fields': "",
    })

    if parameters:
        phantom.act("create ticket", parameters=parameters, assets=['servicenow'], name="create_ticket_3", parent_action=action)    
    else:
        phantom.error("'create_ticket_3' will not be executed due to lack of parameters")
    
    return

def disable_user_1(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('disable_user_1() called')
        
    #phantom.debug('Action: {0} {1}'.format(action['name'], ('SUCCEEDED' if success else 'FAILED')))
    
    # collect data for 'disable_user_1' call
    results_data_1 = phantom.collect2(container=container, datapath=['hunt_file_2:action_result.data.*.process.results.*.username', 'hunt_file_2:action_result.parameter.context.artifact_id'], action_results=results)

    parameters = []
    
    # build parameters list for 'disable_user_1' call
    for results_item_1 in results_data_1:
        if results_item_1[0]:
            parameters.append({
                'username': results_item_1[0],
                # context (artifact id) is added to associate results with the artifact
                'context': {'artifact_id': results_item_1[1]},
            })

    phantom.act(action="disable user", parameters=parameters, assets=['domainctrl1'], callback=join_filter_2, name="disable_user_1", parent_action=action)

    return

def block_hash_3(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('block_hash_3() called')
        
    #phantom.debug('Action: {0} {1}'.format(action['name'], ('SUCCEEDED' if success else 'FAILED')))
    
    # collect data for 'block_hash_3' call
    inputs_data_1 = phantom.collect2(container=container, datapath=['hunt_file_2:artifact:*.cef.fileHash', 'hunt_file_2:artifact:*.id'], action_results=results)

    parameters = []
    
    # build parameters list for 'block_hash_3' call
    for inputs_item_1 in inputs_data_1:
        if inputs_item_1[0]:
            parameters.append({
                'hash': inputs_item_1[0],
                'comment': "",
                # context (artifact id) is added to associate results with the artifact
                'context': {'artifact_id': inputs_item_1[1]},
            })

    phantom.act(action="block hash", parameters=parameters, assets=['carbonblack'], callback=join_filter_2, name="block_hash_3", parent_action=action)

    return

def hunt_file_1(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('hunt_file_1() called')
        
    #phantom.debug('Action: {0} {1}'.format(action['name'], ('SUCCEEDED' if success else 'FAILED')))
    
    # collect data for 'hunt_file_1' call
    filtered_results_data_1 = phantom.collect2(container=container, datapath=["filtered-data:filter_1:condition_1:file_reputation_1:action_result.parameter.hash", "filtered-data:filter_1:condition_1:file_reputation_1:action_result.parameter.context.artifact_id"])

    parameters = []
    
    # build parameters list for 'hunt_file_1' call
    for filtered_results_item_1 in filtered_results_data_1:
        if filtered_results_item_1[0]:
            parameters.append({
                'hash': filtered_results_item_1[0],
                'type': "binary",
                'range': "",
                # context (artifact id) is added to associate results with the artifact
                'context': {'artifact_id': filtered_results_item_1[1]},
            })

    phantom.act(action="hunt file", parameters=parameters, assets=['carbonblack'], callback=block_hash_2, name="hunt_file_1")

    return

def block_hash_2(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('block_hash_2() called')
        
    #phantom.debug('Action: {0} {1}'.format(action['name'], ('SUCCEEDED' if success else 'FAILED')))
    
    # collect data for 'block_hash_2' call
    inputs_data_1 = phantom.collect2(container=container, datapath=['hunt_file_1:artifact:*.cef.fileHash', 'hunt_file_1:artifact:*.id'], action_results=results)

    parameters = []
    
    # build parameters list for 'block_hash_2' call
    for inputs_item_1 in inputs_data_1:
        if inputs_item_1[0]:
            parameters.append({
                'hash': inputs_item_1[0],
                'comment': "",
                # context (artifact id) is added to associate results with the artifact
                'context': {'artifact_id': inputs_item_1[1]},
            })

    phantom.act(action="block hash", parameters=parameters, assets=['carbonblack'], callback=filter_3, name="block_hash_2", parent_action=action)

    return

def file_reputation_1(action=None, success=None, container=None, results=None, handle=None, filtered_artifacts=None, filtered_results=None, custom_function=None, **kwargs):
    phantom.debug('file_reputation_1() called')

    # collect data for 'file_reputation_1' call
    container_data = phantom.collect2(container=container, datapath=['artifact:*.cef.fileHash', 'artifact:*.id'])

    parameters = []
    
    # build parameters list for 'file_reputation_1' call
    for container_item in container_data:
        if container_item[0]:
            parameters.append({
                'hash': container_item[0],
                # context (artifact id) is added to associate results with the artifact
                'context': {'artifact_id': container_item[1]},
            })

    phantom.act(action="file reputation", parameters=parameters, assets=['virustotal'], callback=filter_1, name="file_reputation_1")

    return

def on_finish(container, summary):
    phantom.debug('on_finish() called')
    # This function is called after all actions are completed.
    # summary of all the action and/or all details of actions
    # can be collected here.

    # summary_json = phantom.get_summary()
    # if 'result' in summary_json:
        # for action_result in summary_json['result']:
            # if 'action_run_id' in action_result:
                # action_results = phantom.get_action_results(action_run_id=action_result['action_run_id'], result_data=False, flatten=False)
                # phantom.debug(action_results)

    return