diff --git a/openrelik_api_client/files.py b/openrelik_api_client/files.py index 4e11611..8ee99e2 100644 --- a/openrelik_api_client/files.py +++ b/openrelik_api_client/files.py @@ -84,12 +84,14 @@ def get_file_content( else: raise ValueError("Invalid return_type. Must be 'bytes' or 'text'.") - def download_file(self, file_id: int, filename: str) -> str | None: + def download_file(self, file_id: int, filename: str, output_dir: str = None) -> 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. + output_dir: The directory to save the downloaded file. If None, a temporary + directory will be used. Returns: str: The path to the downloaded file. @@ -99,7 +101,7 @@ def download_file(self, file_id: int, filename: str) -> str | None: response.raise_for_status() filename_prefix, extension = os.path.splitext(filename) file = tempfile.NamedTemporaryFile( - mode="wb", prefix=f"{filename_prefix}", suffix=extension, delete=False + mode="wb", prefix=f"{filename_prefix}", suffix=extension, delete=False, dir=output_dir ) file.write(response.content) file.close() diff --git a/openrelik_api_client/folders.py b/openrelik_api_client/folders.py index d1eb282..a163512 100644 --- a/openrelik_api_client/folders.py +++ b/openrelik_api_client/folders.py @@ -63,6 +63,23 @@ def list_folder(self, folder_id: int) -> list[dict[str, Any]]: response.raise_for_status() return response.json() + def list_subfolders(self, folder_id: int) -> list[dict[str, Any]]: + """List subfolders in a folder. + + Args: + folder_id: The ID of the folder to check. + + Returns: + A list of dictionaries containing subfolder metadata + + Raises: + HTTPError: If the API request failed. + """ + endpoint = f"{self.api_client.base_url}/folders/{folder_id}/folders/" + 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. diff --git a/openrelik_api_client/workflows.py b/openrelik_api_client/workflows.py index e001e74..d694a9e 100644 --- a/openrelik_api_client/workflows.py +++ b/openrelik_api_client/workflows.py @@ -176,3 +176,18 @@ def get_workflow_report(self, workflow_id: int) -> dict[str, Any] | None: response.raise_for_status() if response.status_code == 200: return response.json() + + def get_workflow_templates(self) -> dict[str, Any] | None: + """Retrieves all available workflow templates. + + Returns: + List of workflow templates. + + Raises: + HTTPError: If the API request failed. + """ + endpoint = f"{self.api_client.base_url}/workflows/templates/" + response = self.api_client.session.get(endpoint) + response.raise_for_status() + if response.status_code == 200: + return response.json() diff --git a/tests/files_api_test.py b/tests/files_api_test.py index 19e35d7..7f3a77e 100644 --- a/tests/files_api_test.py +++ b/tests/files_api_test.py @@ -80,7 +80,7 @@ def test_download_file_with_extension(self, mock_temp_file, mock_splitext): mock_splitext.assert_called_once_with("test_file.txt") mock_response.raise_for_status.assert_called_once() mock_temp_file.assert_called_once_with( - mode="wb", prefix="test_file", suffix=".txt", delete=False + mode="wb", prefix="test_file", suffix=".txt", delete=False, dir=None ) mock_file.write.assert_called_once_with(b"file content") assert result == "/tmp/test_file.txt" diff --git a/tests/folders_api_test.py b/tests/folders_api_test.py index 8edd47a..29e1adf 100644 --- a/tests/folders_api_test.py +++ b/tests/folders_api_test.py @@ -103,6 +103,26 @@ def test_list_folder(self): f"{self.api_client.base_url}/folders/123/files/" ) + def test_list_subfolders(self): + """Test list_subfolders method.""" + # Setup + mock_response = MagicMock() + mock_response.json.return_value = [ + {"id": 10, "name": "Subfolder1"}, + {"id": 20, "name": "Subfolder2"}, + ] + self.api_client.session.get.return_value = mock_response + + # Execute + result = self.folders_api.list_subfolders(123) + + # Verify + mock_response.raise_for_status.assert_called_once() + assert result == [{"id": 10, "name": "Subfolder1"}, {"id": 20, "name": "Subfolder2"}] + self.api_client.session.get.assert_called_once_with( + f"{self.api_client.base_url}/folders/123/folders/" + ) + def test_create_root_folder(self): """Test create_root_folder method.""" # Setup diff --git a/tests/workflows_api_test.py b/tests/workflows_api_test.py index 49e40fd..19de880 100644 --- a/tests/workflows_api_test.py +++ b/tests/workflows_api_test.py @@ -62,3 +62,22 @@ def test_get_workflow_report_no_content(self): mock_response.raise_for_status.assert_called_once() mock_response.json.assert_not_called() assert report is None + + def test_get_workflow_templates(self): + """Test get_workflow_templates method.""" + mock_response = MagicMock() + mock_response.status_code = 200 + expected_templates = [ + {"id": 1, "name": "Template 1"}, + {"id": 2, "name": "Template 2"}, + ] + mock_response.json.return_value = expected_templates + self.api_client.session.get.return_value = mock_response + + templates = self.workflows_api.get_workflow_templates() + + self.api_client.session.get.assert_called_once_with( + f"{self.api_client.base_url}/workflows/templates/" + ) + mock_response.raise_for_status.assert_called_once() + assert templates == expected_templates