import json
import logging
import re
import sys
from pathlib import Path

import pandas as pd
from tqdm import tqdm


class SafeDict(dict):
    def __missing__(self, key):
        return "{" + key + "}"


class DetectMaliciousLogs:
    def __init__(self, dumps_dir, ioc_dir):

        self.dumps_dir = Path(dumps_dir)
        self.ioc_dir = Path(ioc_dir)
        self.detection_RISK = 2
        self.suspicion_RISK = 1
        self.threshold = 5
        self.results = {}

    def _rename_if_required(self, registry_path):

        if "HKCU" in registry_path[:4]:
            registry_path = registry_path.replace("HKCU", "HKEY_CURRENT_USER", 1)

        if "HKLM" in registry_path[:4]:
            registry_path = registry_path.replace("HKLM", "HKEY_LOCAL_MACHINE", 1)

        if "HKU" in registry_path[:4]:
            registry_path = registry_path.replace("HKU", "HKEY_USERS", 1)

        if "HKCC" in registry_path[:4]:
            registry_path = registry_path.replace("HKCC", "HKEY_CURRENT_CONFIG", 1)

        if "HKCR" in registry_path[:4]:
            registry_path = registry_path.replace("HKCR", "HKEY_CLASSES_ROOT", 1)

        return registry_path

    def DetectMaliciousIPs(self):
        logging.info("Finding IOC of IPs")
        file_name = self.dumps_dir / "network_changes.csv"

        if not file_name.exists():
            logging.error(f"{file_name} not found")
            return []

        with open(file_name, "r") as malicious_ips_file:
            malicious_ips = malicious_ips_file.read().split("\n")

        data = pd.read_csv(self.dumps_dir / "network_changes.csv")
        df = pd.DataFrame(data)

        ips_contacted = (df.ip_src.append(df.ip_dest)).unique()
        results = [ip for ip in ips_contacted if ip in malicious_ips]

        return results

    def DetectMaliciousProcess(self):
        logging.info("Finding IOC of Process names")
        events_file = self.dumps_dir / "events.csv"

        if not events_file.exists():
            logging.error(f"{events_file} not found")
            return [], [], [], []

        process_data = pd.read_csv(events_file)
        df = pd.DataFrame(process_data)
        process_list = df.Process_Name.unique()

        with open(self.ioc_dir / "process_names.json", "r") as process_names_file:
            process_names = json.loads(process_names_file.read())

        results = []
        detections = []
        for process in process_list:

            if process in process_names["Trojans"]:
                results.append(process)
                detections.append("Trojan")

            if process in process_names["Adwares"]:
                results.append(process)
                detections.append("Adware")

            if process in process_names["Backdoors"]:
                results.append(process)
                detections.append("Backdoor")

            if process in process_names["Downloaders"]:
                results.append(process)
                detections.append("Downloader")

            if process in process_names["Diallers"]:
                results.append(process)
                detections.append("Dialler")

        return results, process_names["system_process"], process_list, detections

    def DetectMaliciousFileAccess(self):
        logging.info("Finding IOC of File Access")

        process_dir = self.dumps_dir / "process"
        process_files = process_dir.rglob("*")

        with open(
            self.ioc_dir / "suspicious_file_locations.json"
        ) as malicious_file_locations_file:
            malicious_file_locations = json.loads(malicious_file_locations_file.read())

        results = {}
        for process_file in tqdm(process_files):
            with open(process_dir / process_file.name, "r") as api_log:
                data = json.loads(api_log.read())
                apicalls = data["Events"]
                process = data["Process-Name"]

            results[process] = {}
            results[process]["DETECTIONS"] = set()
            results[process]["SUSPICIONS"] = set()

            for apicall in apicalls:
                if "File" in apicall[0] or "Directory" in apicall[0]:
                    for malicious_file_location in malicious_file_locations[
                        "DETECTIONS"
                    ]:
                        if re.compile(
                            malicious_file_location.strip()
                            .format(username="[a-zA-Z0-9]*", location="[^ \\\\]*$")
                            .encode("unicode-escape")
                            .decode()
                        ).search(apicall[1]):
                            results[process]["DETECTIONS"].add(
                                malicious_file_location.strip()
                            )

                    for malicious_file_location in malicious_file_locations[
                        "SUSPICIONS"
                    ]:
                        if re.compile(
                            malicious_file_location.strip()
                            .format(username="[a-zA-Z0-9]*", location="[^ \\\\]*$")
                            .encode("unicode-escape")
                            .decode()
                        ).search(apicall[1]):
                            results[process]["SUSPICIONS"].add(
                                malicious_file_location.strip()
                            )

            results[process]["DETECTIONS"] = list(results[process]["DETECTIONS"])
            results[process]["SUSPICIONS"] = list(results[process]["SUSPICIONS"])

        return results

    def DetectMaliciousRegistryAccess(self):
        logging.info("Finding IOC of Registry Access")

        process_dir = self.dumps_dir / "process"
        process_files = process_dir.rglob("*")

        with open(
            self.ioc_dir / "registry-keys-used-by-malware.json"
        ) as malicious_registry_keys_file:
            malicious_registry_keys = json.loads(malicious_registry_keys_file.read())

        results = {}
        for process_file in tqdm(process_files):
            with open(process_dir / process_file.name, "r") as api_log:
                data = json.loads(api_log.read())
                apicalls = data["Events"]
                process = data["Process-Name"]

                results[process] = {}
                results[process]["DETECTIONS"] = set()
                results[process]["SUSPICIONS"] = set()

                for apicall in apicalls:
                    if "Reg" in apicall[0]:
                        for malicious_registry_key in malicious_registry_keys[
                            "DETECTIONS"
                        ]:
                            if re.compile(
                                malicious_registry_key.strip()
                                .format_map(SafeDict(key="[a-zA-Z0-9]*"))
                                .encode("unicode-escape")
                                .decode()
                            ).match(self._rename_if_required(apicall[1])):
                                results[data["Process-Name"]]["DETECTIONS"].add(
                                    apicall[1]
                                )

                        for malicious_registry_key in malicious_registry_keys[
                            "SUSPICIONS"
                        ]:
                            if re.compile(
                                malicious_registry_key.strip()
                                .format_map(SafeDict(key="[a-zA-Z0-9]*"))
                                .encode("unicode-escape")
                                .decode()
                            ).match(self._rename_if_required(apicall[1])):
                                results[data["Process-Name"]]["SUSPICIONS"].add(
                                    apicall[1]
                                )

                results[process]["DETECTIONS"] = list(results[process]["DETECTIONS"])
                results[process]["SUSPICIONS"] = list(results[process]["SUSPICIONS"])

        return results

    def run(self, json_file, required_file):

        high_risk_ip = self.DetectMaliciousIPs()
        (
            malicious_process,
            system_process,
            total_process,
            detections,
        ) = self.DetectMaliciousProcess()
        files_results = self.DetectMaliciousFileAccess()
        registry_results = self.DetectMaliciousRegistryAccess()
        results = {}
        for process in total_process:
            if process not in system_process:
                results[process] = {}
                results[process]["RISK"] = 0
                results[process]["DETECTIONS"] = []
                results[process]["SUSPICIONS"] = []

        for process, detection in zip(malicious_process, detections):
            results[process]["RISK"] = 5
            results[process]["DETECTIONS"].append(detection)

        for process in results.keys():
            results[process]["DETECTIONS"] = (
                results[process]["DETECTIONS"]
                + files_results[process]["DETECTIONS"]
                + registry_results[process]["DETECTIONS"]
            )
            results[process]["SUSPICIONS"] = (
                files_results[process]["SUSPICIONS"]
                + registry_results[process]["SUSPICIONS"]
            )
            results[process]["RISK"] = (
                results[process]["RISK"]
                + (len(results[process]["DETECTIONS"]) * 2)
                + len(results[process]["SUSPICIONS"])
            )

        with open(self.dumps_dir / json_file, "w") as file:
            file.write(json.dumps(results, indent=4))

        df = (
            pd.DataFrame(results)
            .transpose()
            .explode("SUSPICIONS")
            .explode("DETECTIONS")
            .fillna("null")
        )
        df.index.name = "ProcessName"
        df.to_csv(self.dumps_dir / "log_report.csv")

        malicious_process = []

        for process in results.keys():
            if results[process]["RISK"] >= self.threshold:
                malicious_process.append(process)

        process_file = self.dumps_dir / "process.csv"

        if not process_file.exists():
            logging.error("process.csv not found")
            sys.exit(1)

        df = pd.read_csv(process_file)

        full_paths = (
            df[df["ProcessName"].isin(malicious_process)]["ImagePath"].unique().tolist()
        )
        with open(self.dumps_dir / required_file, "w") as file:
            file.write("\n".join(full_paths))
