Expanded API Wrapper

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: