Demisto FSF Client

demisto_fsf_client.py

# Python template - reading arguments, calling a command, handling errors and returning results
res = []
# Constant and mandatory arguments
args = demisto.args()

fileid = args["file"]
orig_path = "demisto"
tags = args.get("tags")


import os
import sys
import socket
import argparse
import struct
import json
import time
import hashlib
import random

from datetime import datetime as dt


class Config:
    HOST = "127.0.0.1"
    PORT = 5800
    LOGFILE = "tmplog.txt"


class FSFClient:
    def __init__(self, config, fullpath, filename, delete, source, archive, suppress_report, full, file):

        self.fullpath = fullpath
        self.filename = filename
        self.delete = delete
        self.source = source
        self.archive = archive
        self.suppress_report = suppress_report
        self.full = full
        self.file = file
        # If multiple server candidates are given, we randomly choose one
        self.host = config.HOST
        self.port = config.PORT
        self.logfile = config.LOGFILE

    # Send files to server for processing and await results
    def process_files(self):

        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        filedata = struct.pack(str(len(self.file)) + "s", self.file)

        msg = '%sFSF_RPC%sFSF_RPC%sFSF_RPC%sFSF_RPC%sFSF_RPC' % (self.filename, self.source, self.archive, self.suppress_report, self.full)
        msg = struct.pack(str(len(msg)) + "s" + str(len(filedata)) + "s", str(msg), filedata)
        buffer = struct.pack('>I', len(msg)) + 'FSF_RPC' + msg

        try:
            sock.connect((self.host, self.port))
            sock.sendall(buffer)
        except Exception as err:
            error = '%s There was a problem sending file %s to %s on port %s. Error: %s\n' % (dt.now(), self.filename, self.host, self.port, str(err))
            print(error)
            sock.close()
            return None

        finally:

            if not self.suppress_report:
                return self.process_results(sock)

            sock.close()
            return None

    # Process the results sent back from the FSF server
    def process_results(self, sock):

        try:
            raw_msg_len = sock.recv(4)
            msg_len = struct.unpack('>I', raw_msg_len)[0]
            data = ''

            while len(data) < msg_len:
                recv_buff = sock.recv(msg_len - len(data))
                data += recv_buff
            sock.close()
            return data

        except Exception as err:

            error = '%s There was a problem getting data for %s from %s on port %s. Error: %s' % (dt.now(), self.filename, self.host, self.port, str(err))
            print(error)
            return None

resultEntries = demisto.executeCommand("getFilePath", {"id": fileid})
while True:
    contents = resultEntries[0].get("Contents")

    if not isinstance(contents, dict):
        print(contents)
        break

    filepath = contents.get("path")
    name = contents.get("name")

    with open(filepath , 'rb') as file:
        data = file.read()

    import time
    #time.sleep(100)

    #from fsfclient import *

    try:
        config = Config()
        config.HOST = demisto.args()["hostname"]
        client = FSFClient(config=config,fullpath=orig_path, filename=name, delete=False, source={}, archive=False, suppress_report=False, full=False, file=data)
    except Exception as err:
        print (err)
        break
    success = False
    for attempt in range(0, 5):
        try:
            results = client.process_files()
            if not results:
                continue
            success = True

        except Exception as err:
            print(err)
            continue

        try:
            import json
            json_results = json.loads(str(results))
            demisto.setContext("FSF_SCAN." + fileid, json_results)
            res = {"ContentsFormat": formats["json"], "Type": entryTypes["note"], "Contents": json_results}
            if tags:
                tags = tags.replace(" ", "")
                tags = tags.split(",")
                res["Tags"] = tags
            demisto.results(res)
            data = []
            data.append({"Key":"Alert", "Value":json_results['Alert']})
            if json_results['Alert']:
                if len(json_results['Summary']['Alert_Yara']) > 0:
                    data.append({"Key":"Alerting YARA Sigs", "Value":json_results['Summary']['Alert_Yara']})
                else:
                    data.append({"Key":"Alerting YARA Sigs", "Value":"None"})
                if len(json_results['Summary']['Alert_JQ']) > 0:
                    data.append({"Key":"Alerting JQ Sigs", "Value":json_results['Summary']['Alert_JQ']})
                else:
                    data.append({"Key":"Alerting JQ Sigs", "Value":"None"})
            if len(json_results['Summary']['Observations']) > 0:
                data.append({"Key":"Observations", "Value":json_results['Summary']['Observations']})
            else:
                data.append({"Key":"Observations", "Value":"None"})
            if len(json_results['Summary']['Yara']) > 0:
                data.append({"Key":"Yara Sigs", "Value":json_results['Summary']['Yara']})
            else:
                data.append({"Key":"Yara Sigs", "Value":"None"})
            if len(json_results['Summary']['Modules']) > 0:
                data.append({"Key":"Modules", "Value":json_results['Summary']['Modules']})
            else:
                data.append({"Key":"Modules", "Value":"None"})

            data = flattenTable(data)
            res = {"ContentsFormat": formats["table"], "Type": entryTypes["note"], "Contents": data}
            if tags:
                res["Tags"] = tags
            demisto.results(res)
            break
        except Exception as err:
            print(str(err))
            success = False
            break
    if not success:
        print("Unable to query FSF with the file")
    break