diff --git a/services/API.py b/services/API.py index 928b1f4..22cfa7d 100644 --- a/services/API.py +++ b/services/API.py @@ -64,6 +64,23 @@ class AirlockAPIWrapper: 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: """ @@ -75,6 +92,12 @@ class AirlockAPIWrapper: 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.""" @@ -115,14 +138,30 @@ class AirlockAPIWrapper: payload = {"groupid": groupid} result = self._post("/v1/agent/find", payload) return pd.DataFrame(result["response"]["agents"]) - - #Baseline Managment + + # Baseline Managment def baseline_find_all(self) -> pd.DataFrame: - """Retrieve all agents.""" - result = self._post("/v1/agent/find", {}) - return pd.DataFrame(result["response"]["agents"]) + """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: