#    Copyright (C) 2025 James Brotosky, Brandon Wickline
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.

import inspect
import json
import logging
import os
import re
from dataclasses import dataclass
from datetime import datetime
from typing import List, Optional

import dotenv
import pandas as pd

from services.policyhandler import pullPolicyExechistories
from utils.utils import colorText, load_env, load_env_json, regulator

logger = logging.getLogger(__name__)

dotenv.load_dotenv()


class Hash:
    """
    Hash model representing Hash data
    """

    def __init__(
        self,
        sha256,
        applications=None,
        baselines=None,
        blocklists=None,
        createtime=None,
        datetime=None,
        description=None,
        filename=None,
        filepath=None,
        filesize=None,
        md5=None,
        modtime=None,
        origname=None,
        productname=None,
        productversion=None,
        publisher=None,
        reputation=None,
        sha128=None,
        sha384=None,
        sha512=None,
    ):
        self.sha256 = sha256
        self.applications = applications
        self.baselines = baselines
        self.blocklists = blocklists
        self.createtime = createtime
        self.datetime = datetime
        self.description = description
        self.filename = filename
        self.filepath = filepath
        self.filesize = filesize
        self.md5 = md5
        self.modtime = modtime
        self.origname = origname
        self.productname = productname
        self.productversion = productversion
        self.publisher = publisher
        self.reputation = reputation
        self.sha128 = sha128
        self.sha384 = sha384
        self.sha512 = sha512

    def __repr__(self):
        attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
        return f"<Hash({attrs})>"

    def __eq__(self, other):
        if isinstance(other, Hash):
            return self.sha256 == other.sha256
        return False

    def __hash__(self):
        return hash(self.sha256)

    def to_dict(self):
        """Returns a dictionary representation of the hash."""
        return self.__dict__

    @staticmethod
    def safe_int(value, default=0):
        """Safely convert a value to int, returning default on failure."""
        try:
            return int(value)
        except (TypeError, ValueError):
            return default


    @classmethod
    def deduplicate(cls, hash_list):
        """
        Deduplicates a list of Hash objects based on sha256.
        Args:
            hash_list (list): List of Hash instances.
        Returns:
            list: Deduplicated list of Hash instances.
        """
        seen = set()
        deduped = []
        for h in hash_list:
            if h.sha256 not in seen:
                seen.add(h.sha256)
                deduped.append(h)
        return deduped

    @classmethod
    def categorize_hashes(cls, hashes):
        import re
        from utils.utils import load_env, load_env_json, regulator

        threat_tolerance = load_env("VT_THREAT_TOLERANCE", cast_type=int)
        bad_publishers_pattern = regulator(load_env_json("BAD_PUBLISHERS", "[]"))
        pups_pattern = regulator(load_env_json("PUPS", "[]"))

        needs_review = []
        approved = []
        unapproved = []

        def reputationtool(hash_obj):
            val = hash_obj.reputation.get("scannermatch") if isinstance(hash_obj.reputation, dict) else None
            if val in [None, "N/A"]:
                return hash_obj.publisher == "Not Signed"
            try:
                return int(val) > threat_tolerance
            except (ValueError, TypeError):
                return hash_obj.publisher == "Not Signed"

        for hash_obj in hashes:
            publisher = hash_obj.publisher or ""
            description = hash_obj.description or ""
            reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
            rep_status = reputation.get("status")

            rep_flag = reputationtool(hash_obj)

            is_signed = publisher != "Not Signed"
            is_untrusted = re.search(bad_publishers_pattern, publisher, re.IGNORECASE) is not None
            is_pup = re.search(pups_pattern, description, re.IGNORECASE) is not None
            has_known_status = rep_status == "KNOWN"

            if (not is_signed and rep_flag) or rep_status == "UNKNOWN":
                needs_review.append(hash_obj)
            elif (
                (is_signed and not is_untrusted and has_known_status and not is_pup) or
                (not is_signed and not rep_flag and not is_untrusted and has_known_status and not is_pup)
            ):
                approved.append(hash_obj)
            else:
                unapproved.append(hash_obj)

        return needs_review, approved, unapproved


    @classmethod
    def export_to_csv(cls, hash_list, directory_path):
        """
        Exports a list of Hash objects to a CSV file in the specified directory.
        The filename is derived from the variable name of the list if possible,
        and includes a timestamp to ensure uniqueness.
        """
        filename = "hashes_export.csv"
        frame = inspect.currentframe()
        if frame is not None and frame.f_back is not None:
            callers_local_vars = frame.f_back.f_locals.items()
            for var_name, var_val in callers_local_vars:
                if var_val is hash_list:
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    filename = f"{var_name}_{timestamp}.csv"
                    break
        else:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"hashes_export_{timestamp}.csv"

        os.makedirs(directory_path, exist_ok=True)
        file_path = os.path.join(directory_path, filename)

        df = pd.DataFrame([h.to_dict() for h in hash_list])
        df.to_csv(file_path, index=False)

        logger.info(f"CSV file saved to: {file_path}")


"""
#Example - Convert Dataframe returned by hash query into hash objects
hash_objects = []
for _, row in df.iterrows():
    try:
        parsed_data = ast.literal_eval(row['data'])
        hash_obj = Hash(sha256=row['sha256'], **parsed_data)
        hash_objects.append(hash_obj)
    except Exception as e:
        print(f"Error parsing row: {e}")

# Display the created Hash objects
for obj in hash_objects:
    print(obj)

    
# Categorize hashes
needs_review, approved, unapproved = Hash.categorize_hashes(
    hashes=hash_objects,
    threat_tolerance=3,
    untrusted_pattern=untrusted_pattern,
    pups_pattern=pups_pattern
)

  # Deduplicate
    deduped_hashes = Hash.deduplicate(hash_list)

# Specify the directory where you want to save the CSV
output_directory = "C:/Users/Brandon/Documents/HashExports"

# Call the export method
Hash.export_to_csv(hashes_for_export, output_directory)


"""


@dataclass
class ExecutionHistoryRecord:
    # Mandatory fields
    username: str
    hostname: str
    netdomain: str
    filename: str
    ppolicy: str
    policyname: str
    policyver: str
    commandline: str
    publisher: str
    sha256: str
    datetime: str

    # Optional fields
    type: Optional[int] = None
    pprocess: Optional[str] = None
    gprocess: Optional[str] = None
    md5: Optional[str] = None
    sha128: Optional[str] = None
    sha384: Optional[str] = None
    sha512: Optional[str] = None
    ip: Optional[str] = None
    localip: Optional[str] = None
    extid: Optional[str] = None
    extname: Optional[str] = None
    exttype: Optional[int] = None  # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension
    extbrowser: Optional[int] = None  # 1 = Chrome, 2 = Firefox, 3 = Edge

    @staticmethod
    def enrich_with_hashes_and_export(
        executions: list, hashes: list, directory_path: str, label: str = "enriched"
    ):
        exec_df = pd.DataFrame([e.__dict__ for e in executions])
        hash_df = pd.DataFrame([h.to_dict() for h in hashes])

        logger.debug(f"Execution DataFrame columns: {exec_df.columns}")
        logger.debug(f"Hash DataFrame columns: {hash_df.columns}")

        if hash_df.empty:
            logger.warning(f"hash_df is empty for label: {label}. Skipping merge.")
            merged_df = exec_df.copy()
        else:
            merged_df = pd.merge(
                exec_df,
                hash_df,
                on="sha256",
                how="left",  # Preserve all executions, enrich where possible
                suffixes=("_exec", "_hash")
            )
            logger.info(f"Merged {len(merged_df)} rows. Non-null hash matches: {merged_df['sha256'].notna().sum()}")

        filename = f"{label}_executions.csv"
        os.makedirs(directory_path, exist_ok=True)
        file_path = os.path.join(directory_path, filename)
        merged_df.to_csv(file_path, index=False)

        logger.info(f"CSV file saved to: {file_path}")


    @classmethod
    def from_dict(cls, data: dict):
        mandatory_fields = [
            "username",
            "hostname",
            "netdomain",
            "filename",
            "ppolicy",
            "policyname",
            "policyver",
            "commandline",
            "publisher",
            "sha256",
            "datetime",
        ]
        missing_fields = [
            field for field in mandatory_fields if field not in data or data[field] is None
        ]
        if missing_fields:
            raise ValueError(f"Missing mandatory fields: {missing_fields}")

        return cls(
            username=data["username"],
            hostname=data["hostname"],
            netdomain=data["netdomain"],
            filename=data["filename"],
            ppolicy=data["ppolicy"],
            policyname=data["policyname"],
            policyver=data["policyver"],
            commandline=data["commandline"],
            publisher=data["publisher"],
            sha256=data["sha256"],
            datetime=data["datetime"],
            type=data.get("type"),
            pprocess=data.get("pprocess"),
            gprocess=data.get("gprocess"),
            md5=data.get("md5"),
            sha128=data.get("sha128"),
            sha384=data.get("sha384"),
            sha512=data.get("sha512"),
            ip=data.get("ip"),
            localip=data.get("localip"),
            extid=data.get("extid"),
            extname=data.get("extname"),
            exttype=data.get("exttype"),
            extbrowser=data.get("extbrowser"),
        )

    @classmethod
    def from_policies(
        cls, api, selected_policies, type_: list, history_days: int
    ) -> List["ExecutionHistoryRecord"]:
        executions = []
        for policy in selected_policies:
            execs = pullPolicyExechistories(
                api, policy, type_, history_days, True
            )
            if execs:
                data = json.loads(execs)
                exechistories = data.get("response", {}).get("exechistories", [])
                if not exechistories:
                    continue

                df = pd.DataFrame(exechistories)
                df = df.drop_duplicates(subset=["sha256", "filename", "hostname"])
                df = df.sort_values(by=["sha256", "filename"])

                executions.extend([cls.from_dict(row.to_dict()) for _, row in df.iterrows()])
                logger.debug(f"Staging of Execution history for policy: {policy.name} is complete")
                print(
                    colorText(
                        f"Staging of Execution history for policy: {policy.name} is complete",
                        "green",
                    )
                )

        return executions

    def __repr__(self):
        attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
        return f"<Execution({attrs})>"


"""
executions = ExecutionHistoryRecord.from_policies(api, selected_policies, type_=[0,1,3], history_days=30)

ExecutionHistoryRecord.enrich_with_hashes_and_export(executions, hash_objects, "C:/Users/Brandon/Documents/EnrichedExports")
"""