openrelik-api-client 0.2.4__tar.gz → 0.4.0__tar.gz

{openrelik_api_client-0.2.4 → openrelik_api_client-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: openrelik-api-client
-Version: 0.2.4
+Version: 0.4.0
 Summary: API client that automatically handles token refresh
 Author: Johan Berggren
 Author-email: jberggren@gmail.com

openrelik_api_client-0.4.0/openrelik_api_client/api_client.py

@@ -0,0 +1,174 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+from typing import Any
+
+import requests
+from requests.exceptions import RequestException
+
+from .configs import ConfigsAPI
+from .files import FilesAPI
+
+
+class APIClient:
+    """
+    A reusable API client that automatically handles token refresh on 401 responses.
+
+    Attributes:
+        api_server_url (str): The URL of the API server.
+        api_key (str): The API key.
+        api_version (str): The API version.
+        session (requests.Session): The session object.
+
+    Example usage:
+        client = APIClient(api_server_url, refresh_token)
+        response = client.get("/users/me/")
+        print(response.json())
+    """
+
+    def __init__(
+        self,
+        api_server_url,
+        api_key=None,
+        api_version="v1",
+    ):
+        self.base_url = f"{api_server_url}/api/{api_version}"
+        self.session = TokenRefreshSession(api_server_url, api_key)
+
+    def get(self, endpoint, **kwargs):
+        """Sends a GET request to the specified API endpoint."""
+        url = f"{self.base_url}{endpoint}"
+        return self.session.get(url, **kwargs)
+
+    def post(self, endpoint, data=None, json=None, **kwargs):
+        """Sends a POST request to the specified API endpoint."""
+        url = f"{self.base_url}{endpoint}"
+        return self.session.post(url, data=data, json=json, **kwargs)
+
+    def put(self, endpoint, data=None, **kwargs):
+        """Sends a PUT request to the specified API endpoint."""
+        url = f"{self.base_url}{endpoint}"
+        return self.session.put(url, data=data, **kwargs)
+
+    def patch(self, endpoint, data=None, json=None, **kwargs):
+        """Sends a PATCH request to the specified API endpoint."""
+        url = f"{self.base_url}{endpoint}"
+        return self.session.patch(url, data=data, json=json, **kwargs)
+
+    def delete(self, endpoint, **kwargs):
+        """Sends a DELETE request to the specified API endpoint."""
+        url = f"{self.base_url}{endpoint}"
+        return self.session.delete(url, **kwargs)
+
+    def get_config(self) -> dict[str, Any]:
+        """
+        DEPRECATED: Use configAPI.get_system_config() instead.
+        This method will be removed in a future version.
+        """
+        warnings.warn(
+            "The 'APIClient.get_config' method is deprecated. Please use 'configAPI.get_system_config()' instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        configs_api = ConfigsAPI(self)
+        return configs_api.get_system_config()
+
+    def download_file(self, file_id: int, filename: str) -> str | None:
+        """
+        DEPRECATED: Use filesAPI.download_file() instead.
+        This method will be removed in a future version.
+        """
+        warnings.warn(
+            "The 'APIClient.download_file' method is deprecated. Please use 'filesAPI.download_file()' instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        files_api = FilesAPI(self)
+        return files_api.download_file(file_id, filename)
+
+    def upload_file(self, file_path: str, folder_id: int) -> int | None:
+        """
+        DEPRECATED: Use filesAPI.download_file() instead.
+        This method will be removed in a future version.
+        """
+        warnings.warn(
+            "The 'APIClient.upload_file' method is deprecated. Please use 'filesAPI.upload_file()' instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        files_api = FilesAPI(self)
+        return files_api.upload_file(file_path, folder_id)
+
+
+class TokenRefreshSession(requests.Session):
+    """Custom session class that handles automatic token refresh on 401 responses."""
+
+    def __init__(self, api_server_url, api_key):
+        """
+        Initializes the TokenRefreshSession with the API server URL and refresh token.
+
+        Args:
+            api_server_url (str): The URL of the API server.
+            refresh_token (str): The refresh token.
+        """
+        super().__init__()
+        self.api_server_url = api_server_url
+        if api_key:
+            self.headers["x-openrelik-refresh-token"] = api_key
+
+    def request(self, method: str, url: str, **kwargs: dict[str, Any]) -> requests.Response:
+        """Intercepts the request to handle token expiration.
+
+        Args:
+            method (str): The HTTP method.
+            url (str): The URL of the request.
+            **kwargs: Additional keyword arguments for the request.
+
+        Returns:
+            Response: The response object.
+
+        Raises:
+            Exception: If the token refresh fails.
+        """
+        response = super().request(method, url, **kwargs)
+
+        if response.status_code == 401:
+            if self._refresh_token(url):
+                # Retry the original request with the new token
+                response = super().request(method, url, **kwargs)
+            else:
+                raise RuntimeError("API key has expired")
+
+        return response
+
+    def _refresh_token(self, requested_url: str) -> bool:
+        """Refreshes the access token using the refresh token."""
+        refresh_url = f"{self.api_server_url}/auth/refresh"
+
+        # If the original URL is the same as the refresh URL, do not attempt to refresh as this
+        # indicates a faulty or expired api key. This prevents an infinite loop of refresh attempts.
+        if requested_url == refresh_url:
+            return False
+
+        try:
+            response = self.get(refresh_url)
+            response.raise_for_status()
+            # Update session headers with the new access token
+            new_access_token = response.json().get("new_access_token")
+            self.headers["x-openrelik-access-token"] = new_access_token
+            return True
+        except RequestException as e:
+            print(f"Failed to refresh token: {e}")
+            return False

openrelik_api_client-0.4.0/openrelik_api_client/configs.py

@@ -0,0 +1,31 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from openrelik_api_client.api_client import APIClient
+
+
+class ConfigsAPI:
+    def __init__(self, api_client: "APIClient"):
+        super().__init__()
+        self.api_client = api_client
+
+    def get_system_config(self) -> dict[str, Any]:
+        """Gets the current OpenRelik system configuration."""
+        endpoint = f"{self.api_client.base_url}/configs/system/"
+        response = self.api_client.session.get(endpoint)
+        response.raise_for_status()
+        return response.json()

openrelik_api_client-0.2.4/openrelik_api_client/api_client.py → openrelik_api_client-0.4.0/openrelik_api_client/files.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,71 +17,73 @@ import os
 import tempfile
 import time
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
-import requests
-from requests.exceptions import RequestException
 from requests_toolbelt import MultipartEncoder
 
+if TYPE_CHECKING:
+    from openrelik_api_client.api_client import APIClient
 
-class APIClient:
-    """
-    A reusable API client that automatically handles token refresh on 401 responses.
-
-    Attributes:
-        api_server_url (str): The URL of the API server.
-        api_key (str): The API key.
-        api_version (str): The API version.
-        session (requests.Session): The session object.
-
-    Example usage:
-        client = APIClient(api_server_url, refresh_token)
-        response = client.get("/users/me/")
-        print(response.json())
-    """
-
-    def __init__(
-        self,
-        api_server_url,
-        api_key=None,
-        api_version="v1",
-    ):
-        self.base_url = f"{api_server_url}/api/{api_version}"
-        self.session = TokenRefreshSession(api_server_url, api_key)
-
-    def get(self, endpoint, **kwargs):
-        """Sends a GET request to the specified API endpoint."""
-        url = f"{self.base_url}{endpoint}"
-        return self.session.get(url, **kwargs)
-
-    def post(self, endpoint, data=None, json=None, **kwargs):
-        """Sends a POST request to the specified API endpoint."""
-        url = f"{self.base_url}{endpoint}"
-        return self.session.post(url, data=data, json=json, **kwargs)
-
-    def put(self, endpoint, data=None, **kwargs):
-        """Sends a PUT request to the specified API endpoint."""
-        url = f"{self.base_url}{endpoint}"
-        return self.session.put(url, data=data, **kwargs)
-
-    def patch(self, endpoint, data=None, json=None, **kwargs):
-        """Sends a PATCH request to the specified API endpoint."""
-        url = f"{self.base_url}{endpoint}"
-        return self.session.patch(url, data=data, json=json, **kwargs)
-
-    def delete(self, endpoint, **kwargs):
-        """Sends a DELETE request to the specified API endpoint."""
-        url = f"{self.base_url}{endpoint}"
-        return self.session.delete(url, **kwargs)
-
-    def get_config(self) -> dict[str, Any]:
-        """Gets the current OpenRelik configuration."""
-        endpoint = f"{self.base_url}/config/system/"
-        response = self.session.get(endpoint)
+
+class FilesAPI:
+    def __init__(self, api_client: "APIClient"):
+        super().__init__()
+        self.api_client = api_client
+
+    def get_file_metadata(self, file_id: int) -> dict[str, Any]:
+        """Reads a file metadata.
+
+        Args:
+            file_id: The ID of the file to get the metadata from.
+
+        Returns:
+            A dictionary containing file metadata.
+
+        Raises:
+            HTTPError: If the API request failed.
+        """
+        endpoint = f"{self.api_client.base_url}/files/{file_id}/"
+        response = self.api_client.session.get(endpoint)
         response.raise_for_status()
         return response.json()
 
+    def get_file_content(
+        self, file_id: int, max_file_size_bytes: int, return_type: str = "bytes"
+    ) -> str | bytes | None:
+        """Download the content of a file.
+
+        Args:
+            file_id: The ID of the file to download.
+            max_file_size_bytes: Maximum file size to download in bytes.
+            return_type: The type of content to return. Can be "bytes" or "text".
+
+        Returns:
+            The content of the file as bytes if return_type is "bytes", or as a string if
+            return_type is "text". Returns None if the file does not exist.
+
+        Raises:
+            RuntimeError: If the file is too large to download.
+            ValueError: If the return_type is not "bytes" or "text".
+        """
+        # Guard against large files as this will read the entire file into memory.
+        # If downloading larger files than 100MB is needed, use the download_file method instead.
+        MAX_FILE_SIZE_GUARD = 100 * 1024 * 1024  # 100 MB
+
+        filesize = self.get_file_metadata(file_id).get("filesize")
+        if filesize > max_file_size_bytes or filesize > MAX_FILE_SIZE_GUARD:
+            raise RuntimeError("File too large to download")
+
+        endpoint = f"{self.api_client.base_url}/files/{file_id}/download_stream"
+        response = self.api_client.session.get(endpoint)
+        response.raise_for_status()
+        if return_type == "text":
+            return response.text
+        elif return_type == "bytes":
+            return response.content
+        else:
+            raise ValueError("Invalid return_type. Must be 'bytes' or 'text'.")
+
     def download_file(self, file_id: int, filename: str) -> str | None:
         """Downloads a file from OpenRelik.
 
@@ -92,8 +94,9 @@ class APIClient:
         Returns:
             str: The path to the downloaded file.
         """
-        endpoint = f"{self.base_url}/files/{file_id}/download"
-        response = self.session.get(endpoint)
+        endpoint = f"{self.api_client.base_url}/files/{file_id}/download"
+        response = self.api_client.session.get(endpoint)
+        response.raise_for_status()
         filename_prefix, extension = os.path.splitext(filename)
         file = tempfile.NamedTemporaryFile(
             mode="wb", prefix=f"{filename_prefix}", suffix=extension, delete=False
@@ -131,7 +134,9 @@ class APIClient:
             raise FileNotFoundError(f"File {file_path} not found.")
 
         if folder_id:
-            response = self.session.get(f"{self.base_url}/folders/{folder_id}")
+            response = self.api_client.session.get(
+                f"{self.api_client.base_url}/folders/{folder_id}"
+            )
             if response.status_code == 404:
                 return file_id
 
@@ -157,8 +162,8 @@ class APIClient:
                         {"file": (file_path.name, chunk, "application/octet-stream")}
                     )
                     headers = {"Content-Type": encoder.content_type}
-                    response = self.session.post(
-                        f"{self.base_url}{endpoint}",
+                    response = self.api_client.session.post(
+                        f"{self.api_client.base_url}{endpoint}",
                         headers=headers,
                         data=encoder.to_string(),
                         params=params,
@@ -182,58 +187,30 @@ class APIClient:
 
         return file_id
 
-
-class TokenRefreshSession(requests.Session):
-    """Custom session class that handles automatic token refresh on 401 responses."""
-
-    def __init__(self, api_server_url, api_key):
-        """
-        Initializes the TokenRefreshSession with the API server URL and refresh token.
+    def get_sql_schemas(self, file_id: int) -> dict[str, Any]:
+        """Retrieve tables and schemas for a supported SQL file.
 
         Args:
-            api_server_url (str): The URL of the API server.
-            refresh_token (str): The refresh token.
+            file_id: The ID of the file to run the query against.
+
+        Returns:
+            A dictionary containing the results.
         """
-        super().__init__()
-        self.api_server_url = api_server_url
-        if api_key:
-            self.headers["x-openrelik-refresh-token"] = api_key
+        endpoint = f"{self.api_client.base_url}/files/{file_id}/sql/schemas/"
+        response = self.api_client.session.get(endpoint)
+        return response.json()
 
-    def request(self, method: str, url: str, **kwargs: dict[str, Any]) -> requests.Response:
-        """Intercepts the request to handle token expiration.
+    def run_sql_query(self, file_id: int, query: str) -> dict[str, Any]:
+        """Runs a SQL query against a supported SQL file.
 
         Args:
-            method (str): The HTTP method.
-            url (str): The URL of the request.
-            **kwargs: Additional keyword arguments for the request.
+            file_id: The ID of the file to run the query against.
+            query: The SQL query to run.
 
         Returns:
-            Response: The response object.
-
-        Raises:
-            Exception: If the token refresh fails.
+            A dictionary containing the query results.
         """
-        response = super().request(method, url, **kwargs)
-
-        if response.status_code == 401:
-            if self._refresh_token():
-                # Retry the original request with the new token
-                response = super().request(method, url, **kwargs)
-            else:
-                raise Exception("Token refresh failed")
-
-        return response
-
-    def _refresh_token(self) -> bool:
-        """Refreshes the access token using the refresh token."""
-        refresh_url = f"{self.api_server_url}/auth/refresh"
-        try:
-            response = self.get(refresh_url)
-            response.raise_for_status()
-            # Update session headers with the new access token
-            new_access_token = response.json().get("new_access_token")
-            self.headers["x-openrelik-access-token"] = new_access_token
-            return True
-        except RequestException as e:
-            print(f"Failed to refresh token: {e}")
-            return False
+        endpoint = f"{self.api_client.base_url}/files/{file_id}/sql/query/"
+        request_body = {"query": query}
+        response = self.api_client.session.post(endpoint, json=request_body)
+        return response.json()

{openrelik_api_client-0.2.4 → openrelik_api_client-0.4.0}/openrelik_api_client/folders.py

@@ -22,6 +22,47 @@ class FoldersAPI:
         super().__init__()
         self.api_client = api_client
 
+    def list_root_folders(
+        self, limit: int, pagination_metadata: bool = False
+    ) -> list[dict[str, Any]]:
+        """List root folders.
+
+        Args:
+            limit: Maximum number of folders to return.
+            pagination_metadata: If True, include pagination metadata in the response.
+
+        Returns:
+            A list of dictionaries containing folder metadata.
+
+        Raises:
+            HTTPError: If the API request failed.
+        """
+        endpoint = f"{self.api_client.base_url}/folders/all/?page_size={limit}"
+        response = self.api_client.session.get(endpoint)
+        response.raise_for_status()
+        if pagination_metadata:
+            folders = response.json()
+        else:
+            folders = response.json().get("folders", [])
+        return folders
+
+    def list_folder(self, folder_id: int) -> list[dict[str, Any]]:
+        """List files in a folder.
+
+        Args:
+            folder_id: The ID of the folder to check.
+
+        Returns:
+            A list of dictionaries containing file metadata
+
+        Raises:
+            HTTPError: If the API request failed.
+        """
+        endpoint = f"{self.api_client.base_url}/folders/{folder_id}/files/"
+        response = self.api_client.session.get(endpoint)
+        response.raise_for_status()
+        return response.json()
+
     def create_root_folder(self, display_name: str) -> int | None:
         """Create a root folder.
 
@@ -56,12 +97,11 @@ class FoldersAPI:
         Raises:
             HTTPError: If the API request failed.
         """
-        # Corrected: Use the passed folder_id in the endpoint URL
         endpoint = f"{self.api_client.base_url}/folders/{folder_id}/folders/"
         data = {"display_name": display_name}
         response = self.api_client.session.post(endpoint, json=data)
         response.raise_for_status()
-        # Corrected: assign to a different variable or use the return directly
+
         new_folder_id = None
         if response.status_code == 201:
             new_folder_id = response.json().get("id")
@@ -84,9 +124,7 @@ class FoldersAPI:
         response.raise_for_status()
         return response.status_code == 200
 
-    def update_folder(
-        self, folder_id: int, folder_data: dict[str, Any]
-    ) -> dict[str, Any] | None:
+    def update_folder(self, folder_id: int, folder_data: dict[str, Any]) -> dict[str, Any] | None:
         """Updates an existing folder.
 
         Args:

{openrelik_api_client-0.2.4 → openrelik_api_client-0.4.0}/openrelik_api_client/workflows.py

@@ -72,6 +72,25 @@ class WorkflowsAPI:
         if response.status_code == 200:
             return response.json()
 
+    def get_workflow_status(self, folder_id: int, workflow_id: int) -> dict[str, Any]:
+        """Retrieves a workflow status by ID.
+
+        Args:
+            folder_id: The ID of the folder where the workflow exists.
+            workflow_id: The ID of the workflow to retrieve.
+
+        Returns:
+            The workflow status data.
+
+         Raises:
+            HTTPError: If the API request failed.
+        """
+        endpoint = f"{self.folders_url}/{folder_id}/workflows/{workflow_id}/status"
+        response = self.api_client.session.get(endpoint)
+        response.raise_for_status()
+        if response.status_code == 200:
+            return response.json()
+
     def update_workflow(
         self, folder_id: int, workflow_id: int, workflow_data: dict
     ) -> dict[str, Any] | None:

{openrelik_api_client-0.2.4 → openrelik_api_client-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "openrelik-api-client"
-version = "0.2.4"
+version = "0.4.0"
 description = "API client that automatically handles token refresh"
 authors = ["Johan Berggren <jberggren@gmail.com>"]
 readme = "README.md"
@@ -10,6 +10,17 @@ python = "^3.10"
 requests = "^2.32.3"
 requests_toolbelt = "^1.0.0"
 
+[tool.poetry.group.dev.dependencies]
+pylint = "^3.1.0"
+pytest = "^8.0.2"
+pytest-mock = "^3.14.0"
+pytest-cov = "^6.0.0"
+
+[tool.pytest.ini_options]
+python_files = "*_test.py"
+python_classes = "Test*"
+python_functions = "test_*"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"