377 lines
14 KiB
Python
377 lines
14 KiB
Python
# 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: Optional[List[str]] = None
|
|
) -> str:
|
|
"""Retrieve execution history logs."""
|
|
payload = {"type": type, "checkpoint": checkpoint}
|
|
if policy is not None:
|
|
payload["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"]
|
|
|
|
def server_logs(self, checkpoint: str | None = None) -> str:
|
|
"""Retrieves Server Activity History Logs."""
|
|
payload = {}
|
|
if checkpoint is not None:
|
|
payload["checkpoint"] = checkpoint
|
|
result = self._post("/v1/logging/svractivities?checkpoint", payload)
|
|
return result["response"]["svractivities"]
|
|
|
|
|
|
"""
|
|
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)
|
|
"""
|