openrelik-api-client 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

openrelik_api_client/api_client.py

@@ -12,17 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import math
-import os
-import tempfile
-import time
-from pathlib import Path
+import warnings
 from typing import Any
-from uuid import uuid4
 
 import requests
 from requests.exceptions import RequestException
-from requests_toolbelt import MultipartEncoder
+
+from .configs import ConfigsAPI
+from .files import FilesAPI
 
 
 class APIClient:
@@ -76,111 +73,43 @@ class APIClient:
         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)
-        response.raise_for_status()
-        return response.json()
+        """
+        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:
-        """Downloads a file from OpenRelik.
-
-        Args:
-            file_id: The ID of the file to download.
-            filename: The name of the file to download.
-
-        Returns:
-            str: The path to the downloaded file.
         """
-        endpoint = f"{self.base_url}/files/{file_id}/download"
-        response = self.session.get(endpoint)
-        filename_prefix, extension = os.path.splitext(filename)
-        file = tempfile.NamedTemporaryFile(
-            mode="wb", prefix=f"{filename_prefix}", suffix=extension, delete=False
+        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,
         )
-        file.write(response.content)
-        file.close()
-        return file.name
+        files_api = FilesAPI(self)
+        return files_api.download_file(file_id, filename)
 
     def upload_file(self, file_path: str, folder_id: int) -> int | None:
-        """Uploads a file to the server.
-
-        Args:
-            file_path: File contents.
-            folder_id: An existing OpenRelik folder identifier.
-
-        Returns:
-            file_id of the uploaded file or None otherwise.
-
-        Raise:
-            FileNotFoundError: if file_path is not found.
         """
-        MAX_CHUNK_RETRIES = 10  # Maximum number of retries for chunk upload
-        CHUNK_RETRY_INTERVAL = 0.5  # seconds
-
-        file_id = None
-        response = None
-        endpoint = "/files/upload"
-        chunk_size = 10 * 1024 * 1024  # 10 MB
-        resumableTotalChunks = 0
-        resumableChunkNumber = 0
-        resumableIdentifier = uuid4().hex
-        file_path = Path(file_path)
-        resumableFilename = file_path.name
-        if not file_path.exists():
-            raise FileNotFoundError(f"File {file_path} not found.")
-
-        if folder_id:
-            response = self.session.get(f"{self.base_url}/folders/{folder_id}")
-            if response.status_code == 404:
-                return file_id
-
-        with open(file_path, "rb") as fh:
-            total_size = Path(file_path).stat().st_size
-            resumableTotalChunks = math.ceil(total_size / chunk_size)
-            while chunk := fh.read(chunk_size):
-                resumableChunkNumber += 1
-                retry_count = 0
-                while retry_count < MAX_CHUNK_RETRIES:
-                    params = {
-                        "resumableRelativePath": resumableFilename,
-                        "resumableTotalSize": total_size,
-                        "resumableCurrentChunkSize": len(chunk),
-                        "resumableChunkSize": chunk_size,
-                        "resumableChunkNumber": resumableChunkNumber,
-                        "resumableTotalChunks": resumableTotalChunks,
-                        "resumableIdentifier": resumableIdentifier,
-                        "resumableFilename": resumableFilename,
-                        "folder_id": folder_id,
-                    }
-                    encoder = MultipartEncoder(
-                        {"file": (file_path.name, chunk, "application/octet-stream")}
-                    )
-                    headers = {"Content-Type": encoder.content_type}
-                    response = self.session.post(
-                        f"{self.base_url}{endpoint}",
-                        headers=headers,
-                        data=encoder.to_string(),
-                        params=params,
-                    )
-                    if response.status_code == 200 or response.status_code == 201:
-                        # Success, move to the next chunk
-                        break
-                    elif response.status_code == 503:
-                        # Server has issue saving the chunk, retry the upload.
-                        retry_count += 1
-                        time.sleep(CHUNK_RETRY_INTERVAL)
-                    elif response.status_code == 429:
-                        # Rate limit exceeded, cancel the upload and raise an error.
-                        raise RuntimeError("Upload failed, maximum retries exceeded")
-                    else:
-                        # Other errors, cancel the upload and raise an error.
-                        raise RuntimeError("Upload failed")
-
-            if response and response.status_code == 201:
-                file_id = response.json().get("id")
-
-        return file_id
+        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):
@@ -216,17 +145,23 @@ class TokenRefreshSession(requests.Session):
         response = super().request(method, url, **kwargs)
 
         if response.status_code == 401:
-            if self._refresh_token():
+            if self._refresh_token(url):
                 # Retry the original request with the new token
                 response = super().request(method, url, **kwargs)
             else:
-                raise Exception("Token refresh failed")
+                raise RuntimeError("API key has expired")
 
         return response
 
-    def _refresh_token(self) -> bool:
+    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()

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/files.py

@@ -0,0 +1,216 @@
+# 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.
+
+import math
+import os
+import tempfile
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+from requests_toolbelt import MultipartEncoder
+
+if TYPE_CHECKING:
+    from openrelik_api_client.api_client import APIClient
+
+
+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.
+
+        Args:
+            file_id: The ID of the file to download.
+            filename: The name of the file to download.
+
+        Returns:
+            str: The path to the downloaded file.
+        """
+        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
+        )
+        file.write(response.content)
+        file.close()
+        return file.name
+
+    def upload_file(self, file_path: str, folder_id: int) -> int | None:
+        """Uploads a file to the server.
+
+        Args:
+            file_path: File contents.
+            folder_id: An existing OpenRelik folder identifier.
+
+        Returns:
+            file_id of the uploaded file or None otherwise.
+
+        Raise:
+            FileNotFoundError: if file_path is not found.
+        """
+        MAX_CHUNK_RETRIES = 10  # Maximum number of retries for chunk upload
+        CHUNK_RETRY_INTERVAL = 0.5  # seconds
+
+        file_id = None
+        response = None
+        endpoint = "/files/upload"
+        chunk_size = 10 * 1024 * 1024  # 10 MB
+        resumableTotalChunks = 0
+        resumableChunkNumber = 0
+        resumableIdentifier = uuid4().hex
+        file_path = Path(file_path)
+        resumableFilename = file_path.name
+        if not file_path.exists():
+            raise FileNotFoundError(f"File {file_path} not found.")
+
+        if 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
+
+        with open(file_path, "rb") as fh:
+            total_size = Path(file_path).stat().st_size
+            resumableTotalChunks = math.ceil(total_size / chunk_size)
+            while chunk := fh.read(chunk_size):
+                resumableChunkNumber += 1
+                retry_count = 0
+                while retry_count < MAX_CHUNK_RETRIES:
+                    params = {
+                        "resumableRelativePath": resumableFilename,
+                        "resumableTotalSize": total_size,
+                        "resumableCurrentChunkSize": len(chunk),
+                        "resumableChunkSize": chunk_size,
+                        "resumableChunkNumber": resumableChunkNumber,
+                        "resumableTotalChunks": resumableTotalChunks,
+                        "resumableIdentifier": resumableIdentifier,
+                        "resumableFilename": resumableFilename,
+                        "folder_id": folder_id,
+                    }
+                    encoder = MultipartEncoder(
+                        {"file": (file_path.name, chunk, "application/octet-stream")}
+                    )
+                    headers = {"Content-Type": encoder.content_type}
+                    response = self.api_client.session.post(
+                        f"{self.api_client.base_url}{endpoint}",
+                        headers=headers,
+                        data=encoder.to_string(),
+                        params=params,
+                    )
+                    if response.status_code == 200 or response.status_code == 201:
+                        # Success, move to the next chunk
+                        break
+                    elif response.status_code == 503:
+                        # Server has issue saving the chunk, retry the upload.
+                        retry_count += 1
+                        time.sleep(CHUNK_RETRY_INTERVAL)
+                    elif response.status_code == 429:
+                        # Rate limit exceeded, cancel the upload and raise an error.
+                        raise RuntimeError("Upload failed, maximum retries exceeded")
+                    else:
+                        # Other errors, cancel the upload and raise an error.
+                        raise RuntimeError("Upload failed")
+
+            if response and response.status_code == 201:
+                file_id = response.json().get("id")
+
+        return file_id
+
+    def get_sql_schemas(self, file_id: int) -> dict[str, Any]:
+        """Retrieve tables and schemas for a supported SQL file.
+
+        Args:
+            file_id: The ID of the file to run the query against.
+
+        Returns:
+            A dictionary containing the results.
+        """
+        endpoint = f"{self.api_client.base_url}/files/{file_id}/sql/schemas/"
+        response = self.api_client.session.get(endpoint)
+        return response.json()
+
+    def run_sql_query(self, file_id: int, query: str) -> dict[str, Any]:
+        """Runs a SQL query against a supported SQL file.
+
+        Args:
+            file_id: The ID of the file to run the query against.
+            query: The SQL query to run.
+
+        Returns:
+            A dictionary containing the query results.
+        """
+        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/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/workflows.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# Copyright 2024-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.
@@ -25,7 +25,7 @@ class WorkflowsAPI:
         self.folders_url = f"{self.api_client.base_url}/folders"
 
     def create_workflow(
-        self, folder_id: int, file_ids: list, template_id: int = None
+        self, folder_id: int, file_ids: list, template_id: int = None, template_params: dict = {}
     ) -> int | None:
         """Creates a new workflow.
 
@@ -46,6 +46,7 @@ class WorkflowsAPI:
             "folder_id": folder_id,
             "file_ids": file_ids,
             "template_id": template_id,
+            "template_params": template_params,
         }
         response = self.api_client.session.post(endpoint, json=data)
         response.raise_for_status()
@@ -157,3 +158,21 @@ class WorkflowsAPI:
         if response.status_code == 200:
             workflow = response.json()
         return workflow
+
+    def get_workflow_report(self, workflow_id: int) -> dict[str, Any] | None:
+        """Retrieves a workflow report.
+
+        Args:
+            workflow_id: The ID of the workflow to retrieve the report from.
+
+        Returns:
+            The workflow report data.
+
+         Raises:
+            HTTPError: If the API request failed.
+        """
+        endpoint = f"{self.api_client.base_url}/workflows/{workflow_id}/report/"
+        response = self.api_client.session.get(endpoint)
+        response.raise_for_status()
+        if response.status_code == 200:
+            return response.json()

{openrelik_api_client-0.3.0.dist-info → openrelik_api_client-0.5.0.dist-info}/METADATA

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

openrelik_api_client-0.5.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+openrelik_api_client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+openrelik_api_client/api_client.py,sha256=rRRMAnEYSO2Gsy7SFrCA2a4LmpIvqo7uXYDleWqWk7Y,6380
+openrelik_api_client/configs.py,sha256=ns0AF6i9vZq8M-1EKEF4xZ-IpcvDPbTFxn1WMAe_PnU,1127
+openrelik_api_client/files.py,sha256=0lWUsMIjUGsz9CGQJFchcfq3myBSaxSZCKPJ1oyMgWQ,8444
+openrelik_api_client/folders.py,sha256=6Hix2G4WoRPBMQu7HTeV_ZS_lZTj8pZW57NW80riZp4,8118
+openrelik_api_client/groups.py,sha256=zhR2oq9_PG0ZmFgVHf62T7KWeg7h6mJmXOABScOJGyU,6116
+openrelik_api_client/workflows.py,sha256=BwIBSjS0krsJWglPiFdAAtcwRMqNmtWDyz9H7UDgHWE,6187
+openrelik_api_client-0.5.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+openrelik_api_client-0.5.0.dist-info/METADATA,sha256=pcSwYPyaRjtqxonmu6mMPr6YpnKdw4JDkerxXSnYgw8,2114
+openrelik_api_client-0.5.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+openrelik_api_client-0.5.0.dist-info/RECORD,,

openrelik_api_client-0.3.0.dist-info/RECORD

@@ -1,9 +0,0 @@
-openrelik_api_client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-openrelik_api_client/api_client.py,sha256=MjkiGJuU_t9S_n0aQSwPPRpgo1NpjVViC7WFx-PCikg,9130
-openrelik_api_client/folders.py,sha256=RtL94DOSP72to_kLDF9HP94zQfXKui8xoLT-OSdHSxw,6943
-openrelik_api_client/groups.py,sha256=zhR2oq9_PG0ZmFgVHf62T7KWeg7h6mJmXOABScOJGyU,6116
-openrelik_api_client/workflows.py,sha256=5DjpNuqGfbLz-glWPqqIZpL6wtHtt2DWsqG-EpFDeRw,5510
-openrelik_api_client-0.3.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-openrelik_api_client-0.3.0.dist-info/METADATA,sha256=2OMDYJtxbMFRWnDyM_bhUrz-GHDxqDvDVYenzfFMHCw,2114
-openrelik_api_client-0.3.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-openrelik_api_client-0.3.0.dist-info/RECORD,,