

import os
import sys
import ast
import pprint
import re

# import astpretty


class DictToObject(object):

    def __init__(self, dictionary):
        def _traverse(key, element):
            if isinstance(element, dict):
                return key, DictToObject(element)
            else:
                return key, element
        objd = dict(_traverse(k, v) for k, v in dictionary.items())
        self.__dict__.update(objd)


"""
class objectdict(dict):
    def __getattr__(self, name):
        if name in self:
            return self[name]
        else:
            raise AttributeError("No such attribute: " + name)

    def __setattr__(self, name, value):
        self[name] = value

    def __delattr__(self, name):
        if name in self:
            del self[name]
        else:
            raise AttributeError("No such attribute: " + name)


class objectview(object):
    def __init__(self, d):
        self.__dict__ = d
"""


class QueryTransformer(ast.NodeTransformer):
    allowedFunctions = ["len", "any", "int", "float",
                        "abs", "complex", "divmod", "pow", "round", "hasattr"]

    def visit_Call(self, node):
        if not node.func.id in self.allowedFunctions:
            raise Exception("Invalid function requested, unable to continue")
        return node

    def visit_Assign(self, node):
        if len(node.targets) > 1:
            raise Exception(
                "Parse failed, more than one target assignemnt detected")
        """
        # this doesnt work
        if isinstance(node.targets[0], ast.Attribute):
            node.targets[0].ctx = ast.Load()
        if isinstance(node.targets[0], ast.Name):
            node.targets[0].ctx = ast.Load()

        node.targets[0].ctx = ast.Load()
        """
        left = ast.Attribute(
            value=ast.Name(id=node.targets[0].value.id, ctx=ast.Load()),
            attr=node.targets[0].attr,
            ctx=Load()
        )

        # return ast.Expr(value=ast.Compare(left=node.targets[0], ops=[ast.Eq()], comparators=[node.value]))
        return ast.Expr(value=ast.Compare(left=left, ops=[ast.Eq()], comparators=[node.value]))


class QueryInterpreter(object):
    def __init__(self, query):

        # replace = with ==
        self.query = re.sub('=', '==', query.strip())
        while '===' in self.query:
            self.query = self.query.replace('===', '==')
        while '!==' in self.query:
            self.query = self.query.replace('!==', '!=')

        try:
            self.ast = ast.parse(self.query, mode='eval')
            self.ast = QueryTransformer().visit(self.ast)
            # astpretty.pprint(self.ast, indent='    ', show_offsets=False)
            self.clause = compile(self.ast, '<AST>', 'eval')
        except Exception as e:
            raise e

    def compare(self, data):

        evt = DictToObject(data['event'])
        # evt = objectdict(data['event'])
        # artifacts = DictToObject(data)
        artifacts = []
        if 'artifacts' in data:
            for artifact in data['artifacts']:
                artifacts.append(DictToObject(artifact))
        try:
            return eval(self.clause, {"event": evt, "artifacts": artifacts})
        except AttributeError as e:
            # this object attribute does not exist
            return False

    def compareAll(self, data):
        # compare
        ret = []
        for x in data:
            ret.append(self.compare(x))

        return ret


if __name__ == '__main__':

    if len(sys.argv) < 2:
        # print("Usage: %s [string]" % sys.argv[0])
        # sys.exit(0)
        ast_string = "event.app = \"dns\" and ( event.ip_src != event.ip_dst )"
    else:
        ast_string = sys.argv[1]

    # print(ast_string)

    # tree = ast.parse(ast_string)
    # print(ast.dump(tree))
    # tree = QueryTransformer().visit(tree)
    # print(ast.dump(tree))

    interpreter = QueryInterpreter(ast_string)
    # astpretty.pprint(interpreter.ast, indent='    ', show_offsets=False)
    print(ast.dump(interpreter.ast))

    example_record = {
        "event": {
            "app": "dns",
            "ip_src": "10.0.0.1",
            "ip_dst": "10.0.0.5",
            "target_username": None
        },
        "artifacts": [
            {"type": "file_scan", "malicious": True,
                "score": 100.0, "filename": "Effervescence.exe"},
            {"type": "file_scan", "malicious": True,
                "score": 100.0, "filename": "Lichtenstein32.dll"}
        ]
    }

    print(interpreter.compare(example_record))

    example_record2 = {
        "event": {
            "app": "dns",
            "ip_src": "10.0.0.1",
            "ip_dst": "10.0.0.1",
        },
        "artifacts": [
            {"type": "file_scan", "malicious": True,
                "score": 100.0, "filename": "Effervescence.exe"},
            {"type": "file_scan", "malicious": True,
                "score": 100.0, "filename": "Lichtenstein32.dll"}
        ]
    }

    print(interpreter.compare(example_record2))

    example_record3 = {
        "event": {
            "key": "AS400",
            "hid": "P1001",
            "ip_proto": 6,
            "ip_dport": 514,
            "payload": "2020-08-03 18:05:14 hawkqa-sccm.hawkio.local CEF:0|HAWK|vTTAC|v1.6.3|3006|DNS query is called for the name HAWKQA-SCCM, type 1, query options 140738562121728, Server List , isNetwork query 0, network index 0, interface index 0, is asynchronous query 0|4|source=\"DNS-Client\" QueryName=\"HAWKQA-SCCM\" QueryType=\"1\" QueryOptions=\"140738562121728\" ServerList=\"\" IsNetworkQuery=\"0\" NetworkQueryIndex=\"0\" InterfaceIndex=\"0\" IsAsyncQuery=\"0\" enrichedPID=true QueryResults=\"fe80::f4f1:5d4:8d9e:32a8;::ffff:10.14.0.105;\" QueryStatus=\"0\" enrichedSession=true severity=\"notice\" priority=4 pid=1068 level=29 sys_version=0 category=0 op=\"Info\" Microsoft-Windows-DNS-Client/Operational=true key=0x8000000000000000 id=2401968254 app=\"CcmExec.exe\" tid=3336 channel=\"Microsoft-Windows-DNS-Client/Operational\" pid_user=\"NT AUTHORITY\\SYSTEM\" sid=\"S-1-5-18\" account type=\"User\"",
            "resource_address_orig": "10.14.0.105",
            "ip_sport": 49746,
            "reception_time": "2020-08-03 18:01:33",
            "class_name": "Microsoft Windows",
            "vendor_name": "IBM",
            "resource_id": "134",
            "bayesian_weight": 552,
            "ip_dst_distinct": 1,
            "product_name": "System-i",
            "resource_address6": "",
            "compliance_asset": False,
            "c81f504f-ff09-07c7-6dbf-ec54a4573b0e:Malware/Virus Propagating to Multiple Host (2 in 30m):alert_name:System i HAWK Agent Default Message:": 8,
            "c103f62e-bbf4-1cea-42a4-04c872f8da16:Abnormal Amount of Destination Ports Detected:ip_src:10.14.0.105:": 2,
            "resource_details": "",
            "resource_asset_criticality": 5,
            "date": "2020-08-03 18:01:33",
            "hash": "4C26EB79E0C9ADD7FC0AC8D972DF51AA",
            "engine": "HAWK-ECE-01",
            "ts": "2020-08-03 18:01:33",
            "_id": "5f2850ff67633217000001ba",
            "g_n": ".",
            "b": False,
            "ip_s": "10.14.0.105",
            "ip_d": "10.14.0.200",
            "c": 12,
            "w": 6,
            "a_name": "System i HAWK Agent Default Message",
            "pri": 4,
            "c_key": "Windows",
            "at_name": "Miscellaneous Information",
            "ip_sh": "HAWKQA-SCCM",
            "ip_dh": "hawkqa52-01",
            "rs_n": "Detected Internal to Internal Event, Detected Internal to Internal Event, Detected Internal to Internal Event, Bayesian Signature (Extremely High), Asset Risk - Very Low, Priority 4, Bayesian Signature (Extremely High), Asset Risk - Very Low, Priority 4, Priority 4",
            "r_name": "HAWKQA-SCCM",
            "o_name": "Microsoft Windows",
            "rs_v": "+ 0.0, + 0.0, + 0.0, + 3.50, - 2.0, + 1.0, + 3.50, - 2.0, + 1.0, + 1.0",
            "s_g_cc2": "INT",
            "s_g_r": "Internal",
            "d_g_cc2": "INT",
            "d_g_r": "Internal",
            "d_g_n": "Internal",
            "d_g_c": "Internal",
            "s_g_n": "Internal",
            "s_g_c": "Internal",
            "r_addr": "10.14.0.105",
            "test_array": ["a", "b", "c", "d"]
        },
        "artifacts": [
            {"type": "file_scan", "malicious": True,
                "score": 100.0, "filename": "Effervescence.exe"},
            {"type": "file_scan", "malicious": True,
                "score": 100.0, "filename": "Lichtenstein32.dll"}
        ]
    }

    print(interpreter.compare(example_record3))
