#    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 json
import logging
from typing import Dict, List, Optional

import pandas as pd
import requests

logger = logging.getLogger(__name__)


class AirlockAPIWrapper:
    """
    A wrapper class for interacting with the Airlock API.
    Provides methods for managing agents, policies, hashes, OTPs, and execution history.
    """

    def __init__(self, base_url: str, api_key: str):
        """
        Initialize the API wrapper.

        Parameters:
        - base_url (str): Base URL of the Airlock API.
        - api_key (str): API key for authentication.
        """
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.headers = {"X-APIKey": self.api_key}

    def _post(self, endpoint: str, payload: Optional[dict] = None) -> dict:
        """
        Internal method to send POST requests to the API.

        Parameters:
        - endpoint (str): API endpoint.
        - payload (dict, optional): Request payload.

        Returns:
        - dict: JSON response from the API.
        """
        url = f"{self.base_url}{endpoint}"
        data = json.dumps(payload or {})
        try:
            logger.debug(f"POST Request to {url} with payload: {payload}")
            response = requests.post(url, headers=self.headers, data=data, verify=False)
            response.raise_for_status()
            logger.debug(f"Response received from {url}")
            return response.json()
        except requests.exceptions.RequestException as e:
            logger.error(f"API request failed: {e}")
            raise

    def _post_raw(self, endpoint: str, payload: Optional[dict] = None) -> bytes:
        """
        Send POST request and return raw response content (bytes).
        Useful for XML endpoints.
        """
        url = f"{self.base_url}{endpoint}"
        data = json.dumps(payload or {})
        try:
            logger.debug(f"POST Request to {url} with payload: {payload}")
            response = requests.post(url, headers=self.headers, data=data, verify=False)
            response.raise_for_status()
            logger.debug(f"Raw response received from {url}")
            return response.content  # bytes
        except requests.exceptions.RequestException as e:
            logger.error(f"API request failed: {e}")
            raise

    # Allowlist Management
    def allowlist_find_all(self) -> pd.DataFrame:
        """
        Retrieve all applications in the allowlist.

        Returns:
        - pd.DataFrame: DataFrame containing allowlisted applications.
        """
        result = self._post("/v1/application", {})
        return pd.DataFrame(result["response"]["applications"])

    def allowlist_export(self, applicationid) -> bytes:
        """Return Allowlist XML as bytes to save to file"""
        payload = {"applicationid": applicationid}
        result = self._post_raw("/v1/application/export", payload)
        return result  # should be bytes

    # Agent Management
    def agent_find_all(self) -> pd.DataFrame:
        """Retrieve all agents."""
        result = self._post("/v1/agent/find", {})
        return pd.DataFrame(result["response"]["agents"])

    def agent_find_by_hostname(self, hostname: str) -> pd.DataFrame:
        """Find agents by hostname."""
        payload = {"hostname": hostname}
        result = self._post("/v1/agent/find", payload)
        return pd.DataFrame(result["response"]["agents"])

    def agent_find_by_id(self, agentid: str) -> pd.DataFrame:
        """Find agents by agent ID."""
        payload = {"agentid": agentid}
        result = self._post("/v1/agent/find", payload)
        return pd.DataFrame(result["response"]["agents"])

    def agent_find_by_status(self, status: int) -> pd.DataFrame:
        """Find agents by status (0 = Offline, 1 = Online, 3 = Safemode)."""
        payload = {"status": status}
        result = self._post("/v1/agent/find", payload)
        return pd.DataFrame(result["response"]["agents"])

    def agent_find_by_username(self, username: str) -> pd.DataFrame:
        """Find agents by username."""
        payload = {"username": username}
        result = self._post("/v1/agent/find", payload)
        return pd.DataFrame(result["response"]["agents"])

    def agent_move(self, agentid: str, groupid: str) -> dict:
        """Move an agent to a different group."""
        payload = {"agentid": agentid, "groupid": groupid}
        return self._post("/v1/agent/move", payload)

    def agents_find_by_group(self, groupid: str) -> pd.DataFrame:
        """Find agents by group ID."""
        payload = {"groupid": groupid}
        result = self._post("/v1/agent/find", payload)
        return pd.DataFrame(result["response"]["agents"])

    # Baseline Managment
    def baseline_find_all(self) -> pd.DataFrame:
        """Retrieve all Baselines."""
        result = self._post("/v1/baseline", {})
        return pd.DataFrame(result["response"]["baselines"])

    def baseline_export(self, baselineid) -> bytes:
        """Return Baseline XML as bytes to save to file"""
        payload = {"baselineid": baselineid}
        result = self._post_raw("/v1/baseline/export", payload)
        return result

    # Blocklist Managment
    def blocklist_find_all(self) -> pd.DataFrame:
        """Retrieve all Baselines."""
        result = self._post("/v1/blocklist", {})
        return pd.DataFrame(result["response"]["blocklists"])

    def blocklist_export(self, blocklistid) -> bytes:
        """Return Blocklist XML as bytes to save to file"""
        payload = {"blocklistid": blocklistid}
        result = self._post_raw("/v1/blocklist/export", payload)
        return result

    # Hash Management
    def hash_add_to_allowlist(self, applicationid: str, hashes: List[str]) -> dict:
        """Add hashes to the allowlist for a specific application."""
        payload = {"applicationid": applicationid, "hashes": hashes}
        return self._post("/v1/hash/application/add", payload)

    def hash_query(self, hashes: List[str]) -> pd.DataFrame:
        """Query information about specific hashes."""
        payload = {"hashes": hashes}
        result = self._post("/v1/hash/query", payload)
        return pd.DataFrame(result["response"]["results"])

    # OTP Management
    def otp_find_active(self) -> pd.DataFrame:
        """Find active OTPs."""
        payload = {"status": "1"}
        result = self._post("/v1/otp/usage", payload)
        return pd.DataFrame(result["response"]["otpusage"])

    def otp_find_awaiting(self) -> pd.DataFrame:
        """Find OTPs that are awaiting activation."""
        payload = {"status": "0"}
        result = self._post("/v1/otp/usage", payload)
        return pd.DataFrame(result["response"]["otpusage"])

    def otp_find_enforced(self) -> pd.DataFrame:
        """Find OTPs that are awaiting activation."""
        payload = {"status": "2"}
        result = self._post("/v1/otp/usage", payload)
        return pd.DataFrame(result["response"]["otpusage"])

    def otp_find_revoked(self) -> pd.DataFrame:
        """Find OTPs that are awaiting activation."""
        payload = {"status": "3"}
        result = self._post("/v1/otp/usage", payload)
        return pd.DataFrame(result["response"]["otpusage"])

    def otp_find_by_agent(self, agentid) -> pd.DataFrame:
        """Find OTP by agent."""
        payload = {"agentid": agentid}
        result = self._post("/v1/otp/usage", payload)
        return pd.DataFrame(result["response"]["otpusage"])

    def otp_generate(self, agentid: str, duration: int, purpose: str) -> str:
        """Generate a new OTP for an agent."""
        payload = {
            "duration": str(duration),
            "agentid": str(agentid),
            "purpose": purpose,
        }
        result = self._post("/v1/otp/retrieve", payload)
        return result["response"]["otpcode"]

    def otp_get_activities(self, otpid: str) -> pd.DataFrame:
        """Retrieve activities associated with a specific OTP."""
        payload = {"otpid": otpid}
        result = self._post("/v1/otp/activities", payload)
        return pd.DataFrame(result["response"]["otpactivities"])

    def otp_revoke(self, otpid: str) -> dict:
        """
        Revoke an active OTP.
        Parameters:
        - otpid (str): The ID of the OTP to revoke.
        Returns:
        - dict: JSON response from the API.
        """
        payload = {"otpid": otpid}
        return self._post("/v1/otp/revoke", payload)

    def otp_validate(self, otpcode: str) -> dict:
        """
        Validate an OTP code.
        Parameters:
        - otpcode (str): The OTP code to validate.
        Returns:
        - dict: JSON response indicating validity.
        """
        payload = {"otpcode": otpcode}
        return self._post("/v1/otp/validate", payload)

    # Policy Management
    def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict:
        """Add path exclusions to a policy group."""
        payload = {"groupid": groupid, "path": paths}
        return self._post("/v1/group/path/add", payload)

    def policy_add_publishers(self, groupid: str, publishers: List[str]) -> dict:
        """Add publishers to a policy group."""
        payload = {"groupid": groupid, "publisher": publishers}
        return self._post("/v1/group/publisher/add", payload)

    def policy_clone(self, source_groupid: str, target_groupid: str) -> dict:
        """Clone a policy from one group to another."""
        payload = {"groupid": source_groupid, "targetgroupid": target_groupid}
        return self._post("/v1/group/assign", payload)

    def policy_find_all(self) -> pd.DataFrame:
        """Retrieve all policy groups."""
        result = self._post("/v1/group")
        return pd.DataFrame(result["response"]["groups"])

    def policy_list_agents(self, groupid: str) -> pd.DataFrame:
        """List agents assigned to a specific policy group."""
        payload = {"groupid": groupid}
        result = self._post("/v1/group/agents", payload)
        return pd.DataFrame(result["response"]["agents"])

    def policy_list_allowlists(self, groupid: str) -> pd.DataFrame:
        """List allowlists assigned to a specific policy group."""
        payload = {"groupid": groupid}
        result = self._post("/v1/group/policies", payload)
        return pd.DataFrame(result["response"]["applications"])

    def policy_set_auditmode(self, groupid: str, auditmode: str) -> dict:
        """Set audit mode for a policy group. 1=Audit, 0=Enforcement"""
        payload = {"groupid": groupid, "auditmode": auditmode}
        return self._post("/v1/group/settings/auditmode", payload)

    def policy_set_script_custom(
        self,
        groupid: str,
        script_custom: int,
        scripts_audit: List[str],
        scripts_disabled: List[str],
        scripts_respect: List[str],
    ) -> dict:
        """Set audit mode for a policy group. 1=Audit, 0=Enforcement"""
        payload = {
            "groupid": groupid,
            "script_custom": script_custom,
            "scripts_audit": scripts_audit,
            "scripts_disabled": scripts_disabled,
            "scripts_respect": scripts_respect,
        }
        return self._post("/v1/group/settings/script_custom", payload)

    def policy_set_upgradetarget(
        self,
        groupid: str,
        windows: str,
        macos: str,
    ) -> dict:
        payload = {
            "groupid": groupid,
            "windows": windows,
            "macos": macos,
        }
        return self._post("/v1/group/settings/selfupgrade/target", payload)

    # Execution History
    def history_logging(
        self, type: List[str], checkpoint: str, policy: List[str]
    ) -> str:
        """Retrieve execution history logs."""
        payload = {"type": type, "checkpoint": checkpoint, "policy": policy}
        result = self._post("/v1/logging/exechistories", payload)
        return result["response"]["exechistories"]

    def history_execution(
        self, today: str, date_selected: str, agent_name: str
    ) -> List[Dict]:
        """
        Retrieve execution history logs.

        "datefrom":"", //(Optional) Datefrom is for date range search, formatted as "YYYY-MM-DD"
        "dateto":"", //(Optional) Dateto is for date range search, formatted as "YYYY-MM-DD"
        "category":"", //(Optional) Category for filtering type
        "hostname":"", //(Optional) Hostname to filter
        "username":"admin", //(Optional) Username to filter
        "netdomain":"", //(Optional) Domain (or group) to filter
        "filename":"", //(Optional) Filename to filter
        "ppolicy":"", //(Optional) Parent Policy name to filter
        "policyname":"", //(Optional) Policy name to filter
        "policyver":"", //(Optional) Policy version to filter (e.g. "v95")
        "commandline":"", //(Optional) Commandline to filter
        "publisher":"", //(Optional) Publisher to filter
        "pprocess":"", //(Optional) Parent Process to filter
        "sha256":"", //(Optional) SHA256 hash to filter
        "contains":["hostname"], //(Optional) Contains is an array for wildcard searches on a filter
        "limit":"5" //(Optional) Limit the amount of results returned, default set to 50

        """

        payload = {"datefrom": date_selected, "dateto": today, "hostname": agent_name}
        result = self._post("/v1/getexechistory", payload)
        return result["response"]["exechistory"]


"""
from services.API import AirlockAPIWrapper


api = AirlockAPIWrapper(base_url="https://airlock.example.com/api", api_key="your_api_key_here")

#Example: Get all agents

agents_df = api.agent_find_all()
print("All Agents:")
print(agents_df)
"""
