# 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 . 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 # 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"]) # 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"]) # 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_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"]) # 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) # 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) """