ThreatX

import demistomock as demisto
from CommonServerPython import *

from CommonServerUserPython import *

""" IMPORTS """
import socket
import struct
import time
from operator import itemgetter

import requests
import urllib3

# disable insecure warnings
urllib3.disable_warnings()

""" GLOBAL VARS """
CUSTOMER_NAME = ""
API_KEY = ""
BASE_URL = ""
DBOT_THRESHOLD = 70
USE_SSL = True

""" HELPER FUNCTIONS """


def http_request(url_suffix, commands=None):
    state = demisto.getIntegrationContext()

    session_token = state.get("session_token")

    if url_suffix != "/login":
        demisto.info(f"running request with url={BASE_URL + url_suffix} with commands={commands}")
        data = {"token": session_token, "customer_name": CUSTOMER_NAME}
    else:
        demisto.info(f"running request with url={BASE_URL + url_suffix}")
        data = {}

    if commands is not None:
        data.update(commands)

    res = requests.post(BASE_URL + url_suffix, verify=USE_SSL, json=data)

    if res.status_code != requests.codes.ok:
        if url_suffix == "/login":
            demisto.setIntegrationContext({"session_token": None, "token_expires": None})
            demisto.info(f"{res.status_code} from server during login. Clearing session token cache.")

        raise DemistoException(f"HTTP {res.status_code} Error in API call to ThreatX service - {res.text}")

    resp_json = {}  # type:dict
    try:
        resp_json = res.json()
    except ValueError:
        raise DemistoException(f"Could not parse the response from ThreatX: {res.text}")

    if "Ok" not in resp_json:
        if url_suffix == "/login":
            demisto.setIntegrationContext({"session_token": None, "token_expires": None})
            raise DemistoException(f"Login response error - {res.text}.")

        raise DemistoException(res.text)

    if url_suffix == "/login" and "status" in resp_json["Ok"] and resp_json["Ok"]["status"] is not True:
        demisto.setIntegrationContext({"session_token": None, "token_expires": None})
        raise DemistoException("Invalid credentials.")

    return resp_json["Ok"]


@logger
def initialize():
    endpoint = "/login"
    commands = {"command": "login", "api_token": API_KEY}

    state = demisto.getIntegrationContext()

    if not state.get("session_token"):
        session_token = None
        token_expires = None
    else:
        session_token = state.get("session_token")
        token_expires = state.get("token_expires")

    demisto.info("Initializing request...")

    if session_token is None or (token_expires is not None and token_expires < int(time.time())):
        if session_token is None:
            demisto.info("Session token missing - getting new session token...")
        elif token_expires is not None and token_expires < int(time.time()):
            demisto.info("Session token expired - getting new session token...")

        r = http_request(endpoint, commands)
        demisto.setIntegrationContext({"session_token": r["token"], "token_expires": int(time.time() + (10 * 60))})
        return

    demisto.info("Cached session token not expired.")
    return


def pretty_ip(decimal_ip):
    """Convert decimal ip to dotted quad format"""
    packed_ip = struct.pack("!I", decimal_ip)
    return socket.inet_ntoa(packed_ip)


def pretty_time(input_time):
    """Convert unix epoch time to human readable format"""
    return time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime(input_time))


def set_dbot_score(threatx_score):
    """Set the DBot Score based on the ThreatX risk score"""
    if threatx_score >= DBOT_THRESHOLD:
        return 3
    elif threatx_score > 10:
        return 2
    else:
        return 0


""" FUNCTIONS """


@logger
def block_ip(ip, description):
    commands = {"command": "new_blocklist", "entry": {"ip": ip, "description": description, "created": int(time.time())}}

    return http_request("/lists", commands)


@logger
def block_ip_command(args):
    ips = args.get("ip", [])
    description = args.get("description", "Added by ThreatX Demisto Integration")
    results = []
    errors = []
    for ip in argToList(ips):
        try:
            ip_result = block_ip(ip, description)
        except Exception as error:
            demisto.error(f"failed block ip: {ip}\n{traceback.format_exc()}")
            errors.append(f"Failed to block ip: {ip} error: {error}")
        else:
            results.append(ip_result)
    if results:
        readable_outputs = tableToMarkdown("Block IP", results, ["Result"], removeNull=True)

        return_results(
            CommandResults(
                outputs=results, readable_output=readable_outputs, outputs_prefix="IP(val.Address === obj.Address).Address"
            )
        )
    if errors:
        return_error("\n".join(errors))


@logger
def unblock_ip(ip):
    commands = {"command": "delete_blocklist", "ip": ip}

    return http_request("/lists", commands)


@logger
def unblock_ip_command(args):
    ip = args.get("ip", None)
    results = unblock_ip(ip)

    md = tableToMarkdown("Unblock IP", results, ["Result"], removeNull=True)

    ec = {"IP(val.Address === obj.Address)": {"Address": ip}}

    return_outputs(md, ec, results)


@logger
def blacklist_ip(ip, description):
    commands = {"command": "new_blacklist", "entry": {"ip": ip, "description": description, "created": int(time.time())}}

    return http_request("/lists", commands)


@logger
def blacklist_ip_command(args):
    ips = args.get("ip", None)
    description = args.get("description", "Added by ThreatX Demisto Integration")
    results = []
    errors = []
    for ip in argToList(ips):
        try:
            ip_result = blacklist_ip(ip, description)
        except Exception as error:
            demisto.error(f"failed adding ip: {ip} to balcklist\n{traceback.format_exc()}")
            errors.append(f"Failed to add ip: {ip} to blacklist error: {error}")
        else:
            results.append(ip_result)

    if results:
        readable_outputs = tableToMarkdown("Blacklist IP", results, ["Result"], removeNull=True)

        return_results(
            CommandResults(
                outputs=results, readable_output=readable_outputs, outputs_prefix="IP(val.Address === obj.Address).Address"
            )
        )
    if errors:
        return_error("\n".join(errors))


@logger
def unblacklist_ip(ip):
    commands = {"command": "delete_blacklist", "ip": ip}

    return http_request("/lists", commands)


@logger
def unblacklist_ip_command(args):
    ip = args.get("ip", None)
    results = unblacklist_ip(ip)

    md = tableToMarkdown("Unblacklist IP", results, ["Result"], removeNull=True)

    ec = {"IP(val.Address === obj.Address)": {"Address": ip}}

    return_outputs(md, ec, results)


@logger
def whitelist_ip(ip, description):
    commands = {"command": "new_whitelist", "entry": {"ip": ip, "description": description, "created": int(time.time())}}

    return http_request("/lists", commands)


@logger
def whitelist_ip_command(args):
    ip = args.get("ip", None)
    description = args.get("description", "Added by ThreatX Demisto Integration")
    results = whitelist_ip(ip, description)

    md = tableToMarkdown("Whitelist IP", results, ["Result"], removeNull=True)

    ec = {"IP(val.Address === obj.Address)": {"Address": ip}}

    return_outputs(md, ec, results)


@logger
def unwhitelist_ip(ip):
    commands = {"command": "delete_whitelist", "ip": ip}

    return http_request("/lists", commands)


@logger
def unwhitelist_ip_command(args):
    ip = args.get("ip", None)
    results = unwhitelist_ip(ip)

    md = tableToMarkdown("Unwhitelist IP", results, ["Result"], removeNull=True)

    ec = {"IP(val.Address === obj.Address)": {"Address": ip}}

    return_outputs(md, ec, results)


@logger
def get_entities(entity_name, entity_id, entity_ip, timeframe):
    commands = {"command": "list", "query": {}}  # type: dict

    if entity_name is not None:
        entity_names = entity_name.split(",")
        my_entity_name = {"codenames": entity_names}
        commands["query"].update(my_entity_name)

    if entity_id is not None:
        entity_ids = entity_id.split(",")
        my_entity_id = {"entity_ids": entity_ids}
        commands["query"].update(my_entity_id)

    if entity_ip is not None:
        entity_ips = entity_ip.split(",")
        my_entity_ip = {"ip_addresses": entity_ips}
        commands["query"].update(my_entity_ip)

    first_seen = None

    if timeframe is not None:
        if timeframe == "1-Hour":
            first_seen = int(time.time() - (60 * 60))
        elif timeframe == "1-Day":
            first_seen = int(time.time() - (24 * 60 * 60))
        elif timeframe == "1-Week":
            first_seen = int(time.time() - (7 * 24 * 60 * 60))
        elif timeframe == "1-Month":
            first_seen = int(time.time() - (31 * 24 * 60 * 60))

    if first_seen:
        my_timeframe = {"first_seen": first_seen}
        commands["query"].update(my_timeframe)

    return http_request("/entities", commands)


@logger
def get_entity_risk(entity_id):
    commands = {"command": "risk_changes", "id": entity_id}

    return http_request("/entities", commands)


@logger
def get_entities_command(args):
    entity_name = args.get("entity_name", None)
    entity_id = args.get("entity_id", None)
    entity_ip = args.get("entity_ip", None)
    timeframe = args.get("timeframe", None)
    results = get_entities(entity_name, entity_id, entity_ip, timeframe)

    dbot_scores = []
    ip_enrich = []
    human_readable = []
    entities_context = []
    for entity in results:
        risk_score = 0
        e_risk = None
        # Grab the entity risk so we can set the Dbot score for the Actor IPs
        e_id = entity.get("id")

        if e_id:
            e_risk = get_entity_risk(e_id)

        if isinstance(e_risk, list) and e_risk and isinstance(e_risk[-1], dict) and "risk" in e_risk[-1]:
            risk_score = e_risk[-1]["risk"]

        entity["risk"] = risk_score

        iplist = []

        for actor in entity.get("actors", []):
            if "ip_address" in actor:
                ipdot = pretty_ip(actor["ip_address"])
                iplist.append(ipdot)
                actor["ip_address"] = ipdot

            if "interval_time_start" in actor:
                actor["interval_time_start"] = pretty_time(actor["interval_time_start"])

            if "interval_time_stop" in actor:
                actor["interval_time_stop"] = pretty_time(actor["interval_time_stop"])

            if "fingerprint" in actor and actor.get("fingerprint") is not None and "last_seen" in actor.get("fingerprint", {}):
                actor["fingerprint"]["last_seen"] = pretty_time(actor["fingerprint"]["last_seen"])

            dbscore = set_dbot_score(risk_score)

            dbot_scores.append({"Vendor": "ThreatX", "Indicator": ipdot, "Type": "ip", "Score": dbscore})

            if dbscore == 3:
                ip_enrich.append(
                    {
                        "Address": ipdot,
                        "Malicious": {"Vendor": "ThreatX", "Description": "ThreatX risk score is " + str(risk_score)},
                    }
                )
            else:
                ip_enrich.append({"Address": ipdot})
        entities_context.append({"ID": e_id, "Name": entity["codename"], "IP": iplist, "Risk": risk_score})

        human_readable.append(
            {"Name": entity["codename"], "ID": e_id, "IP Addresses": ", ".join(iplist), "ThreatX Risk Score": risk_score}
        )

    ec = {
        "Threatx.Entity(val.ID && val.ID === obj.ID)": entities_context,
        "DBotScore": dbot_scores,
        "IP(val.Address === obj.Address)": ip_enrich,
    }

    return_outputs(tableToMarkdown("Entities", human_readable), ec, results)


@logger
def get_entity_notes(entity_id):
    commands = {"command": "notes", "id": entity_id}

    return http_request("/entities", commands)


@logger
def get_entity_notes_command(args):
    entity_id = args.get("entity_id", None)
    results = get_entity_notes(entity_id)

    # Reverse sort the list by timestamp
    sorted_results = sorted(results, key=itemgetter("timestamp"), reverse=True)

    # Replace dates with pretty format
    for note in sorted_results:
        if "timestamp" in note:
            note["timestamp"] = pretty_time(note["timestamp"])

    md = tableToMarkdown("Entity Notes", sorted_results, headerTransform=string_to_table_header)

    ec = {"Threatx.Entity(val.ID && val.ID === obj.ID)": {"ID": entity_id, "Note": sorted_results}}

    return_outputs(md, ec, sorted_results)


@logger
def add_entity_note(entity_id, message):
    commands = {"command": "new_note", "note": {"entity_id": entity_id, "content": message}}

    return http_request("/entities", commands)


@logger
def add_entity_note_command(args):
    entity_id = args.get("entity_id", None)
    message = args.get("message", None)
    results = add_entity_note(entity_id, message)

    md = tableToMarkdown("New Entity Note", results, ["Result"], removeNull=True)

    return_outputs(md, {}, results)


@logger
def command_test_module():
    results = http_request("/users", {"command": "list"})

    if not isinstance(results, list):
        raise DemistoException("Unrecognized response from ThreatX.")

    if not results:
        raise DemistoException("Empty response from ThreatX.")

    if "username" not in results[0]:
        raise DemistoException("Unexpected response from ThreatX.")

    demisto.results("ok")


""" EXECUTION CODE """


def main():
    global CUSTOMER_NAME, API_KEY, BASE_URL, DBOT_THRESHOLD, USE_SSL
    params = demisto.params()

    url = params.get("url", "").strip("/")
    BASE_URL = url + "/tx_api/v1"

    CUSTOMER_NAME = params.get("credentials", {}).get("identifier") or params.get("customer_name")
    API_KEY = params.get("credentials", {}).get("password") or params.get("api_key")
    DBOT_THRESHOLD = int(params.get("dbot_threshold", 70))
    USE_SSL = not params.get("insecure")

    command = demisto.command()
    demisto.info(f"command is {command}")
    args = demisto.args()
    try:
        handle_proxy()
        initialize()
        if command == "test-module":
            command_test_module()
        elif command == "threatx-block-ip":
            block_ip_command(args)
        elif command == "threatx-unblock-ip":
            unblock_ip_command(args)
        elif command == "threatx-blacklist-ip":
            blacklist_ip_command(args)
        elif command == "threatx-unblacklist-ip":
            unblacklist_ip_command(args)
        elif command == "threatx-whitelist-ip":
            whitelist_ip_command(args)
        elif command == "threatx-unwhitelist-ip":
            unwhitelist_ip_command(args)
        elif command == "threatx-get-entities":
            get_entities_command(args)
        elif command == "threatx-get-entity-notes":
            get_entity_notes_command(args)
        elif command == "threatx-add-entity-note":
            add_entity_note_command(args)

    except Exception as e:
        return_error(str(e), error=traceback.format_exc())


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