175 lines
7.5 KiB
Python
175 lines
7.5 KiB
Python
import json
|
|
import logging
|
|
from typing import Dict, List, Optional
|
|
|
|
import httpx
|
|
import pandas as pd
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class AirlockAPIWrapper:
|
|
def __init__(self, base_url: str, api_key: str):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.api_key = api_key
|
|
self.headers = {"X-APIKey": self.api_key}
|
|
|
|
async def _post(self, endpoint: str, payload: Optional[dict] = None) -> dict:
|
|
url = f"{self.base_url}{endpoint}"
|
|
data = json.dumps(payload or {})
|
|
timeout = httpx.Timeout(300.0)
|
|
async with httpx.AsyncClient(verify=False, timeout=timeout) as client:
|
|
try:
|
|
logger.debug(f"POST Request to {url} with payload: {payload}")
|
|
response = await client.post(url, headers=self.headers, data=data) # pyright: ignore[reportArgumentType]
|
|
response.raise_for_status()
|
|
logger.debug(f"Response received from {url}")
|
|
return response.json()
|
|
except httpx.RequestError as e:
|
|
logger.error(f"API request failed: {e}")
|
|
raise
|
|
|
|
# Allowlist Management
|
|
async def allowlist_find_all(self) -> pd.DataFrame:
|
|
result = await self._post("/v1/application", {})
|
|
return pd.DataFrame(result["response"]["applications"])
|
|
|
|
# Agent Management
|
|
async def agent_find_all(self) -> pd.DataFrame:
|
|
result = await self._post("/v1/agent/find", {})
|
|
return pd.DataFrame(result["response"]["agents"])
|
|
|
|
async def agent_find_by_hostname(self, hostname: str) -> pd.DataFrame:
|
|
payload = {"hostname": hostname}
|
|
result = await self._post("/v1/agent/find", payload)
|
|
return pd.DataFrame(result["response"]["agents"])
|
|
|
|
async def agent_find_by_id(self, agentid: str) -> pd.DataFrame:
|
|
payload = {"agentid": agentid}
|
|
result = await self._post("/v1/agent/find", payload)
|
|
return pd.DataFrame(result["response"]["agents"])
|
|
|
|
async def agent_find_by_status(self, status: int) -> pd.DataFrame:
|
|
payload = {"status": status}
|
|
result = await self._post("/v1/agent/find", payload)
|
|
return pd.DataFrame(result["response"]["agents"])
|
|
|
|
async def agent_find_by_username(self, username: str) -> pd.DataFrame:
|
|
payload = {"username": username}
|
|
result = await self._post("/v1/agent/find", payload)
|
|
return pd.DataFrame(result["response"]["agents"])
|
|
|
|
async def agent_move(self, agentid: str, groupid: str) -> dict:
|
|
payload = {"agentid": agentid, "groupid": groupid}
|
|
return await self._post("/v1/agent/move", payload)
|
|
|
|
async def agents_find_by_group(self, groupid: str) -> pd.DataFrame:
|
|
payload = {"groupid": groupid}
|
|
result = await self._post("/v1/agent/find", payload)
|
|
return pd.DataFrame(result["response"]["agents"])
|
|
|
|
# Hash Management
|
|
async def hash_add_to_allowlist(self, applicationid: str, hashes: List[str]) -> dict:
|
|
payload = {"applicationid": applicationid, "hashes": hashes}
|
|
return await self._post("/v1/hash/application/add", payload)
|
|
|
|
async def hash_query(self, hashes: List[str]) -> pd.DataFrame:
|
|
payload = {"hashes": hashes}
|
|
result = await self._post("/v1/hash/query", payload)
|
|
return pd.DataFrame(result["response"]["results"])
|
|
|
|
# OTP Management
|
|
async def otp_find_active(self) -> pd.DataFrame:
|
|
payload = {"status": "1"}
|
|
result = await self._post("/v1/otp/usage", payload)
|
|
return pd.DataFrame(result["response"]["otpusage"])
|
|
|
|
async def otp_find_awaiting(self) -> pd.DataFrame:
|
|
payload = {"status": "0"}
|
|
result = await self._post("/v1/otp/usage", payload)
|
|
return pd.DataFrame(result["response"]["otpusage"])
|
|
|
|
async def otp_find_by_agent(self, agentid) -> pd.DataFrame:
|
|
payload = {"agentid": agentid}
|
|
result = await self._post("/v1/otp/usage", payload)
|
|
return pd.DataFrame(result["response"]["otpusage"])
|
|
|
|
async def otp_generate(self, agentid: str, duration: int, purpose: str) -> str:
|
|
payload = {
|
|
"duration": str(duration),
|
|
"agentid": str(agentid),
|
|
"purpose": purpose,
|
|
}
|
|
result = await self._post("/v1/otp/retrieve", payload)
|
|
return result["response"]["otpcode"]
|
|
|
|
async def otp_get_activities(self, otpid: str) -> pd.DataFrame:
|
|
payload = {"otpid": otpid}
|
|
result = await self._post("/v1/otp/activities", payload)
|
|
return pd.DataFrame(result["response"]["otpactivities"])
|
|
|
|
async def otp_revoke(self, otpid: str) -> dict:
|
|
payload = {"otpid": otpid}
|
|
return await self._post("/v1/otp/revoke", payload)
|
|
|
|
async def otp_validate(self, otpcode: str) -> dict:
|
|
payload = {"otpcode": otpcode}
|
|
return await self._post("/v1/otp/validate", payload)
|
|
|
|
# Policy Management
|
|
async def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict:
|
|
payload = {"groupid": groupid, "path": paths}
|
|
return await self._post("/v1/group/path/add", payload)
|
|
|
|
async def policy_add_publishers(self, groupid: str, publishers: List[str]) -> dict:
|
|
payload = {"groupid": groupid, "publisher": publishers}
|
|
return await self._post("/v1/group/publisher/add", payload)
|
|
|
|
async def policy_clone(self, source_groupid: str, target_groupid: str) -> dict:
|
|
payload = {"groupid": source_groupid, "targetgroupid": target_groupid}
|
|
return await self._post("/v1/group/assign", payload)
|
|
|
|
async def policy_find_all(self) -> pd.DataFrame:
|
|
result = await self._post("/v1/group")
|
|
return pd.DataFrame(result["response"]["groups"])
|
|
|
|
async def policy_list_agents(self, groupid: str) -> pd.DataFrame:
|
|
payload = {"groupid": groupid}
|
|
result = await self._post("/v1/group/agents", payload)
|
|
return pd.DataFrame(result["response"]["agents"])
|
|
|
|
async def policy_list_allowlists(self, groupid: str) -> pd.DataFrame:
|
|
payload = {"groupid": groupid}
|
|
result = await self._post("/v1/group/policies", payload)
|
|
return pd.DataFrame(result["response"]["applications"])
|
|
|
|
async def policy_set_auditmode(self, groupid: str, auditmode: str) -> dict:
|
|
payload = {"groupid": groupid, "auditmode": auditmode}
|
|
return await self._post("/v1/group/settings/auditmode", payload)
|
|
|
|
async def policy_set_script_custom(
|
|
self,
|
|
groupid: str,
|
|
script_custom: int,
|
|
scripts_audit: Optional[List[str]] = None,
|
|
scripts_disabled: Optional[List[str]] = None,
|
|
scripts_respect: Optional[List[str]] = None
|
|
) -> dict:
|
|
payload = {
|
|
"groupid": groupid,
|
|
"script_custom": script_custom,
|
|
"scripts_audit": scripts_audit or [],
|
|
"scripts_disabled": scripts_disabled or [],
|
|
"scripts_respect": scripts_respect or []
|
|
}
|
|
return await self._post("/v1/group/settings/script_custom", payload)
|
|
|
|
# Execution History
|
|
async def history_logging(self, type: List[str], checkpoint: str, policy: List[str]) -> List[Dict]:
|
|
payload = {"type": type, "checkpoint": checkpoint, "policy": policy}
|
|
result = await self._post("/v1/logging/exechistories", payload)
|
|
return result["response"]["exechistories"]
|
|
|
|
async def history_execution(self, today: str, date_selected: str, agent_name: str) -> List[Dict]:
|
|
payload = {"datefrom": date_selected, "dateto": today, "hostname": agent_name}
|
|
result = await self._post("/v1/getexechistory", payload)
|
|
return result["response"]["exechistory"] |