AirlockTools/models/execution.py

#    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 dataclasses
from dataclasses import asdict, dataclass
from datetime import datetime
import inspect
import json
import logging
import os
import re
from typing import List, Optional, Tuple

import dotenv
import pandas as pd

import airlock_libs
from services.API import AirlockAPIWrapper
from utils.configmanager import get_system_list, get_system_value
from utils.utils import regulator

logger = logging.getLogger(__name__)

dotenv.load_dotenv()


@dataclass
class Hash:
    """
    Hash model representing Hash data
    """

    sha256: str
    applications: str
    baselines: str
    blocklists: str
    createtime: str
    datetime: str
    description: str
    filename: str
    filepath: str
    filesize: str
    md5: str
    modtime: str
    origname: str
    productname: str
    productversion: str
    publisher: str
    reputation: str
    sha128: str
    sha384: str
    sha512: str
    at_decision: Optional[str] = None

    def to_dict(self):
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict):
        return cls(**data)

    @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):
        threat_tolerance = get_system_value("VT_THREAT_TOLERANCE", cast_type=int)
        bad_publishers_pattern = regulator(get_system_list("BAD_PUBLISHERS"))
        pups_pattern = regulator(get_system_list("PUPS"))

        approved_count = 0
        unapproved_count = 0
        needs_review_count = 0

        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 {}
            )
            scannermatch = reputation.get("scannermatch")

            logger.debug(f"Evaluating hash: {hash_obj}")
            logger.debug(
                f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}"
            )

            # 1. Unapproved: bad publisher or PUP
            if re.search(bad_publishers_pattern, publisher, re.IGNORECASE):
                logger.debug("Unapproved: Publisher matches bad publisher pattern.")
                hash_obj.at_decision = "unapproved"
                unapproved_count += 1
                continue

            if re.search(pups_pattern, description, re.IGNORECASE):
                logger.debug("Unapproved: Description matches PUP pattern.")
                hash_obj.at_decision = "unapproved"
                unapproved_count += 1
                continue

            # 2. Approved: signed
            if publisher != "Not Signed":
                logger.debug("Approved: File is signed and not flagged.")
                hash_obj.at_decision = "approved"
                approved_count += 1
                continue

            # 3. Approved or Unapproved based on threat level
            try:
                score = int(scannermatch)  # pyright: ignore[reportArgumentType]
                logger.debug(f"Parsed scannermatch score: {score}")
                if score > threat_tolerance:  # pyright: ignore[reportOperatorIssue]
                    logger.debug("Unapproved: Unsigned file with high threat score.")
                    hash_obj.at_decision = "unapproved"
                    unapproved_count += 1
                else:
                    logger.debug("Approved: Unsigned file with low threat score.")
                    hash_obj.at_decision = "approved"
                    approved_count += 1
            except (ValueError, TypeError):
                logger.debug(
                    "Needs Review: Scannermatch score is missing or invalid. — {e}"
                )
                hash_obj.at_decision = "needs_review"
                needs_review_count += 1

        logger.debug(
            f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}"
        )
        return hashes

    @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}")


@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
    hash_obj: Optional[Hash] = None

    @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"),
            hash_obj=data.get("hash_obj"),
        )

    @classmethod
    def from_policies(
        cls, api, selected_policies, type_: list, history_days: int
    ) -> List["ExecutionHistoryRecord"]:
        executions = []
        for policy in selected_policies:
            execs = airlock_libs.pull_policy_exec_histories(
                api, policy.name, str([1, 2, 6, 7]), history_days
            )
            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"
                )

        return executions

    @staticmethod
    def enrich_with_hashes(
        api: AirlockAPIWrapper, executions: List["ExecutionHistoryRecord"]
    ) -> List["ExecutionHistoryRecord"]:
        """
        Enriches each ExecutionHistoryRecord with a matching Hash object by querying the API.
        """
        sha256_list = list({e.sha256.strip().lower() for e in executions if e.sha256})
        logger.info(
            f"Extracted {len(sha256_list)} unique sha256 values from {len(executions)} execution records."
        )

        if not sha256_list:
            logger.warning(
                "No sha256 values found in execution records. Skipping enrichment."
            )
            return executions

        logger.debug("Querying hash data from API...")
        hash_df = api.hash_query(sha256_list)
        logger.info(f"Retrieved {len(hash_df)} hash records from API.")

        hash_objects = []
        required_fields = {
            f.name
            for f in dataclasses.fields(Hash)
            if f.default == dataclasses.MISSING
            and f.default_factory == dataclasses.MISSING
        }

        for sha256, (_, row) in zip(sha256_list, hash_df.iterrows()):
            row_dict = row.to_dict()

            # Unwrap nested 'data' field if present
            if "data" in row_dict and isinstance(row_dict["data"], dict):
                row_dict = row_dict["data"]

            # Inject the sha256 back into the row
            row_dict["sha256"] = sha256

            missing = required_fields - row_dict.keys()
            if missing:
                logger.warning(f"Skipping hash row due to missing fields: {missing}")
                logger.debug(f"Row content: {row_dict}")
                continue

            try:
                hash_obj = Hash.from_dict(row_dict)
                hash_objects.append(hash_obj)
            except Exception as e:
                logger.warning(f"Failed to create Hash from row: {e}")
                logger.debug(f"Row content: {row_dict}")

        logger.debug("Converted hash DataFrame to Hash objects.")

        hash_lookup = {h.sha256.strip().lower(): h for h in hash_objects}
        logger.debug("Built hash lookup table.")

        enriched_count = 0
        for exec_record in executions:
            hash_obj = hash_lookup.get(exec_record.sha256.strip().lower())
            if hash_obj:
                exec_record.hash_obj = hash_obj
                enriched_count += 1

        logger.info(
            f"Enriched {enriched_count} out of {len(executions)} execution records with hash data."
        )
        return executions

    @staticmethod
    def categorize_executions_by_hash_decision(
        executions: List["ExecutionHistoryRecord"],
    ) -> List["ExecutionHistoryRecord"]:
        """
        Categorizes the hash_obj of each ExecutionHistoryRecord based on publisher, description, and reputation.

        Modifies the `at_decision` field of each associated Hash object in-place.

        Returns:
            List[ExecutionHistoryRecord]: The same list, with hash_obj.at_decision updated.
        """
        threat_tolerance = get_system_value("VT_THREAT_TOLERANCE", cast_type=int)
        bad_publishers_pattern = regulator(get_system_list("BAD_PUBLISHERS"))
        pups_pattern = regulator(get_system_list("PUPS"))

        approved_count = 0
        unapproved_count = 0
        needs_review_count = 0

        for record in executions:
            hash_obj = record.hash_obj
            if not hash_obj:
                continue  # Skip if no hash object is attached

            publisher = hash_obj.publisher or ""
            description = hash_obj.description or ""
            reputation = (
                hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
            )
            scannermatch = reputation.get("scannermatch")

            logger.debug(f"Evaluating hash: {hash_obj}")
            logger.debug(
                f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}"
            )

            # 1. Unapproved: bad publisher or PUP
            if re.search(bad_publishers_pattern, publisher, re.IGNORECASE):
                logger.debug("Unapproved: Publisher matches bad publisher pattern.")
                hash_obj.at_decision = "unapproved"
                unapproved_count += 1
                continue

            if re.search(pups_pattern, description, re.IGNORECASE):
                logger.debug("Unapproved: Description matches PUP pattern.")
                hash_obj.at_decision = "unapproved"
                unapproved_count += 1
                continue

            # 2. Approved: signed
            if publisher != "Not Signed":
                logger.debug("Approved: File is signed and not flagged.")
                hash_obj.at_decision = "approved"
                approved_count += 1
                continue

            # 3. Approved or Unapproved based on threat level
            try:
                score = int(scannermatch)  # pyright: ignore[reportArgumentType]
                logger.debug(f"Parsed scannermatch score: {score}")
                if threat_tolerance is not None and score >= threat_tolerance:
                    logger.debug("Unapproved: Unsigned file with high threat score.")
                    hash_obj.at_decision = "unapproved"
                    unapproved_count += 1
                else:
                    logger.debug("Approved: Unsigned file with low threat score.")
                    hash_obj.at_decision = "approved"
                    approved_count += 1
            except (ValueError, TypeError) as e:
                logger.debug(
                    f"Needs Review: Scannermatch score is missing or invalid. — {e}"
                )
                hash_obj.at_decision = "needs_review"
                needs_review_count += 1

        logger.debug(
            f"Final counts — Needs Review: {needs_review_count}, "
            f"Approved: {approved_count}, Unapproved: {unapproved_count}"
        )

        return executions

    @classmethod
    def sort_by_hash_decision(cls, executions: List["ExecutionHistoryRecord"]) -> Tuple[
        List["ExecutionHistoryRecord"],
        List["ExecutionHistoryRecord"],
        List["ExecutionHistoryRecord"],
        List["ExecutionHistoryRecord"],
    ]:
        """
        Sorts ExecutionHistoryRecord objects into approved, unapproved, needs_review, and unknown groups
        based on the value of hash_obj.at_decision.

        Returns:
            Tuple of lists: (approved, unapproved, needs_review, unknown)
        """
        approved = []
        unapproved = []
        needs_review = []
        unknown = []

        sorted_executions = sorted(executions, key=lambda x: x.filename)

        for record in sorted_executions:
            decision = getattr(record.hash_obj, "at_decision", None)
            if decision == "approved":
                approved.append(record)
            elif decision == "unapproved":
                unapproved.append(record)
            elif decision == "needs_review":
                needs_review.append(record)
            else:
                unknown.append(record)

        logger.info(
            f"[ExecutionHistoryRecord] Sorted {len(sorted_executions)} records by hash_obj.at_decision:"
        )
        logger.info(f"  Approved: {len(approved)}")
        logger.info(f"  Unapproved: {len(unapproved)}")
        logger.info(f"  Needs Review: {len(needs_review)}")
        logger.info(f"  Unknown/Unset: {len(unknown)}")

        return approved, unapproved, needs_review, unknown


"""
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")
"""