pyattackforge 0.1.3__tar.gz → 0.1.8__tar.gz

{pyattackforge-0.1.3 → pyattackforge-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyattackforge
-Version: 0.1.3
+Version: 0.1.8
 Summary: Python wrapper for the AttackForge API
 Home-page: https://github.com/Tantalum-Labs/PyAttackForge
 Author: Shane S
@@ -40,6 +40,7 @@ A lightweight Python library for interacting with the AttackForge API.
 - Create findings from existing writeups by passing a `writeup_id`
 - Upload evidence to findings or testcases
 - Update/assign testcases to link findings or add notes
+- Link vulnerabilities to testcases via the client
 - Dry-run mode for testing
 
 ---
@@ -112,14 +113,20 @@ client.create_vulnerability(
 
 ### Creating a finding from an existing writeup
 
-If you already have a writeup/library entry and just need to create a finding bound to it, you can either pass `writeup_id` to `create_vulnerability` (as above) or call `create_finding_from_writeup` directly:
+If you already have a writeup/library entry and just need to create a finding bound to it, you can either pass `writeup_id` to `create_vulnerability` (as above) or call `create_finding_from_writeup` directly. Prefer the 24-character writeup id (`id` / `_id`); if only a numeric `reference_id` is available, use that. You can also specify the library key (e.g., `approved_writeups`, `Main Vulnerabilities`):
 
 ```python
 client.create_finding_from_writeup(
     project_id="abc123",
-    writeup_id="68e92c7a821c05c8405a8003",
+    writeup_id="68e92c7a821c05c8405a8003",  # writeup id
+    library="approved_writeups",             # optional: library key/name
     priority="High",
-    affected_assets=[{"name": "ssh-prod-1"}]
+    affected_assets=[{"name": "ssh-prod-1"}],
+    linked_testcases=["5e8017d2e1385f0c58e8f4f8"],  # optional: link testcases at creation
+    likelihood_of_exploitation=5,
+    steps_to_reproduce="1. Do something\n2. Observe result",
+    notes=[{"note": "Created via API", "type": "PLAINTEXT"}],
+    tags=["automation"]
 )
 ```
 
@@ -152,21 +159,49 @@ client.add_note_to_finding(
 
 Add a note/update to a testcase (PUT to the testcase endpoint):
 ```python
-client.update_testcase(
+client.add_note_to_testcase(
     project_id="abc123",
     testcase_id="5e8017d2e1385f0c58e8f4f8",
-    update_fields={
-        "details": "Observed during retest on 2025-09-19."
-    }
+    note="Observed during retest on 2025-09-19.",
+    status="Tested"  # optional
 )
 ```
 
-Associate findings to a testcase (merges with existing linked vulnerabilities if provided):
+Associate findings to a testcase:
 ```python
 client.assign_findings_to_testcase(
     project_id="abc123",
     testcase_id="5e8017d2e1385f0c58e8f4f8",
-    vulnerability_ids=["66849b77950ab45e68fc7b48", "6768d29db1782d7362a2df5f"]
+    vulnerability_ids=["66849b77950ab45e68fc7b48", "6768d29db1782d7362a2df5f"],
+    additional_fields={"status": "Tested"} # optional
+)
+```
+Or link from the vulnerability side using its update endpoint:
+```python
+client.link_vulnerability_to_testcases(
+    vulnerability_id="69273ef0f4a7c85d03930667",
+    testcase_ids=["5e8017d2e1385f0c58e8f4f8"],
+    project_id="abc123",  # optional
+)
+```
+
+Fetch project testcases:
+```python
+testcases = client.get_testcases("abc123")
+```
+
+Fetch a single testcase (if supported in your tenant):
+```python
+testcase = client.get_testcase("abc123", "5e8017d2e1385f0c58e8f4f8")
+```
+
+Merge and add findings to a testcase in one call:
+```python
+client.add_findings_to_testcase(
+    project_id="abc123",
+    testcase_id="5e8017d2e1385f0c58e8f4f8",
+    vulnerability_ids=["69273ef0f4a7c85d03930667"],
+    additional_fields={"status": "Tested"} # optional
 )
 ```
 
@@ -212,32 +247,7 @@ See the source code for full details and docstrings.
 - `create_vulnerability(
       project_id: str,
       title: str,
-      affected_asset_name: str,
-      priority: str,
-      likelihood_of_exploitation: int,
-      description: str,
-      attack_scenario: str,
-      remediation_recommendation: str,
-      steps_to_reproduce: str,
-      tags: Optional[list] = None,
-      notes: Optional[list] = None,
-      is_zeroday: bool = False,
-      is_visible: bool = True,
-      import_to_library: Optional[str] = None,
-      import_source: Optional[str] = None,
-      import_source_id: Optional[str] = None,
-      custom_fields: Optional[list] = None,
-      linked_testcases: Optional[list] = None,
-      custom_tags: Optional[list] = None,
-  ) -> dict`
-
-See the source code for full details and docstrings.
-
----
-- `create_vulnerability(
-      project_id: str,
-      title: str,
-      affected_asset_name: str,
+      affected_assets: list,
       priority: str,
       likelihood_of_exploitation: int,
       description: str,
@@ -254,7 +264,21 @@ See the source code for full details and docstrings.
       custom_fields: Optional[list] = None,
       linked_testcases: Optional[list] = None,
       custom_tags: Optional[list] = None,
+      writeup_custom_fields: Optional[list] = None,
   ) -> dict`
+- `create_finding_from_writeup(project_id: str, writeup_id: str, priority: str, affected_assets: Optional[list] = None, linked_testcases: Optional[list] = None, **kwargs) -> dict`
+- `get_findings_for_project(project_id: str, priority: Optional[str] = None) -> list`
+- `upsert_finding_for_project(...)`
+- `get_vulnerability(vulnerability_id: str) -> dict`
+- `add_note_to_finding(vulnerability_id: str, note: Any, note_type: str = "PLAINTEXT") -> dict`
+- `upload_finding_evidence(vulnerability_id: str, file_path: str) -> dict`
+- `upload_testcase_evidence(project_id: str, testcase_id: str, file_path: str) -> dict`
+- `get_testcases(project_id: str) -> list`
+- `get_testcase(project_id: str, testcase_id: str) -> dict or None`
+- `link_vulnerability_to_testcases(vulnerability_id: str, testcase_ids: List[str], project_id: Optional[str] = None) -> dict`
+- `assign_findings_to_testcase(project_id: str, testcase_id: str, vulnerability_ids: List[str], existing_linked_vulnerabilities: Optional[List[str]] = None, additional_fields: Optional[Dict[str, Any]] = None) -> dict`
+- `add_findings_to_testcase(project_id: str, testcase_id: str, vulnerability_ids: List[str], additional_fields: Optional[Dict[str, Any]] = None) -> dict`
+- `add_note_to_testcase(project_id: str, testcase_id: str, note: str, status: Optional[str] = None) -> dict`
 
 See the source code for full details and docstrings.
 

{pyattackforge-0.1.3 → pyattackforge-0.1.8}/README.md

@@ -12,6 +12,7 @@ A lightweight Python library for interacting with the AttackForge API.
 - Create findings from existing writeups by passing a `writeup_id`
 - Upload evidence to findings or testcases
 - Update/assign testcases to link findings or add notes
+- Link vulnerabilities to testcases via the client
 - Dry-run mode for testing
 
 ---
@@ -84,14 +85,20 @@ client.create_vulnerability(
 
 ### Creating a finding from an existing writeup
 
-If you already have a writeup/library entry and just need to create a finding bound to it, you can either pass `writeup_id` to `create_vulnerability` (as above) or call `create_finding_from_writeup` directly:
+If you already have a writeup/library entry and just need to create a finding bound to it, you can either pass `writeup_id` to `create_vulnerability` (as above) or call `create_finding_from_writeup` directly. Prefer the 24-character writeup id (`id` / `_id`); if only a numeric `reference_id` is available, use that. You can also specify the library key (e.g., `approved_writeups`, `Main Vulnerabilities`):
 
 ```python
 client.create_finding_from_writeup(
     project_id="abc123",
-    writeup_id="68e92c7a821c05c8405a8003",
+    writeup_id="68e92c7a821c05c8405a8003",  # writeup id
+    library="approved_writeups",             # optional: library key/name
     priority="High",
-    affected_assets=[{"name": "ssh-prod-1"}]
+    affected_assets=[{"name": "ssh-prod-1"}],
+    linked_testcases=["5e8017d2e1385f0c58e8f4f8"],  # optional: link testcases at creation
+    likelihood_of_exploitation=5,
+    steps_to_reproduce="1. Do something\n2. Observe result",
+    notes=[{"note": "Created via API", "type": "PLAINTEXT"}],
+    tags=["automation"]
 )
 ```
 
@@ -124,21 +131,49 @@ client.add_note_to_finding(
 
 Add a note/update to a testcase (PUT to the testcase endpoint):
 ```python
-client.update_testcase(
+client.add_note_to_testcase(
     project_id="abc123",
     testcase_id="5e8017d2e1385f0c58e8f4f8",
-    update_fields={
-        "details": "Observed during retest on 2025-09-19."
-    }
+    note="Observed during retest on 2025-09-19.",
+    status="Tested"  # optional
 )
 ```
 
-Associate findings to a testcase (merges with existing linked vulnerabilities if provided):
+Associate findings to a testcase:
 ```python
 client.assign_findings_to_testcase(
     project_id="abc123",
     testcase_id="5e8017d2e1385f0c58e8f4f8",
-    vulnerability_ids=["66849b77950ab45e68fc7b48", "6768d29db1782d7362a2df5f"]
+    vulnerability_ids=["66849b77950ab45e68fc7b48", "6768d29db1782d7362a2df5f"],
+    additional_fields={"status": "Tested"} # optional
+)
+```
+Or link from the vulnerability side using its update endpoint:
+```python
+client.link_vulnerability_to_testcases(
+    vulnerability_id="69273ef0f4a7c85d03930667",
+    testcase_ids=["5e8017d2e1385f0c58e8f4f8"],
+    project_id="abc123",  # optional
+)
+```
+
+Fetch project testcases:
+```python
+testcases = client.get_testcases("abc123")
+```
+
+Fetch a single testcase (if supported in your tenant):
+```python
+testcase = client.get_testcase("abc123", "5e8017d2e1385f0c58e8f4f8")
+```
+
+Merge and add findings to a testcase in one call:
+```python
+client.add_findings_to_testcase(
+    project_id="abc123",
+    testcase_id="5e8017d2e1385f0c58e8f4f8",
+    vulnerability_ids=["69273ef0f4a7c85d03930667"],
+    additional_fields={"status": "Tested"} # optional
 )
 ```
 
@@ -184,32 +219,7 @@ See the source code for full details and docstrings.
 - `create_vulnerability(
       project_id: str,
       title: str,
-      affected_asset_name: str,
-      priority: str,
-      likelihood_of_exploitation: int,
-      description: str,
-      attack_scenario: str,
-      remediation_recommendation: str,
-      steps_to_reproduce: str,
-      tags: Optional[list] = None,
-      notes: Optional[list] = None,
-      is_zeroday: bool = False,
-      is_visible: bool = True,
-      import_to_library: Optional[str] = None,
-      import_source: Optional[str] = None,
-      import_source_id: Optional[str] = None,
-      custom_fields: Optional[list] = None,
-      linked_testcases: Optional[list] = None,
-      custom_tags: Optional[list] = None,
-  ) -> dict`
-
-See the source code for full details and docstrings.
-
----
-- `create_vulnerability(
-      project_id: str,
-      title: str,
-      affected_asset_name: str,
+      affected_assets: list,
       priority: str,
       likelihood_of_exploitation: int,
       description: str,
@@ -226,7 +236,21 @@ See the source code for full details and docstrings.
       custom_fields: Optional[list] = None,
       linked_testcases: Optional[list] = None,
       custom_tags: Optional[list] = None,
+      writeup_custom_fields: Optional[list] = None,
   ) -> dict`
+- `create_finding_from_writeup(project_id: str, writeup_id: str, priority: str, affected_assets: Optional[list] = None, linked_testcases: Optional[list] = None, **kwargs) -> dict`
+- `get_findings_for_project(project_id: str, priority: Optional[str] = None) -> list`
+- `upsert_finding_for_project(...)`
+- `get_vulnerability(vulnerability_id: str) -> dict`
+- `add_note_to_finding(vulnerability_id: str, note: Any, note_type: str = "PLAINTEXT") -> dict`
+- `upload_finding_evidence(vulnerability_id: str, file_path: str) -> dict`
+- `upload_testcase_evidence(project_id: str, testcase_id: str, file_path: str) -> dict`
+- `get_testcases(project_id: str) -> list`
+- `get_testcase(project_id: str, testcase_id: str) -> dict or None`
+- `link_vulnerability_to_testcases(vulnerability_id: str, testcase_ids: List[str], project_id: Optional[str] = None) -> dict`
+- `assign_findings_to_testcase(project_id: str, testcase_id: str, vulnerability_ids: List[str], existing_linked_vulnerabilities: Optional[List[str]] = None, additional_fields: Optional[Dict[str, Any]] = None) -> dict`
+- `add_findings_to_testcase(project_id: str, testcase_id: str, vulnerability_ids: List[str], additional_fields: Optional[Dict[str, Any]] = None) -> dict`
+- `add_note_to_testcase(project_id: str, testcase_id: str, note: str, status: Optional[str] = None) -> dict`
 
 See the source code for full details and docstrings.
 

{pyattackforge-0.1.3 → pyattackforge-0.1.8}/pyattackforge/__init__.py

@@ -1,24 +1,22 @@
-"""
-PyAttackForge
-
-A lightweight Python library for interacting with the AttackForge API.
-"""
-
-"""
-PyAttackForge is free software: you can redistribute it and/or modify
-it under the terms of the GNU Affero General Public License as published by
-the Free Software Foundation, either version 3 of the License, or
-(at your option) any later version.
-
-PyAttackForge is distributed in the hope that it will be useful,
-but WITHOUT ANY WARRANTY; without even the implied warranty of
-MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-GNU Affero General Public License for more details.
-
-You should have received a copy of the GNU Affero General Public License
-along with this program.  If not, see <https://www.gnu.org/licenses/>.
-"""
-
-from .client import PyAttackForgeClient
-
-__all__ = ["PyAttackForgeClient"]
+"""
+PyAttackForge
+
+A lightweight Python library for interacting with the AttackForge API.
+
+PyAttackForge is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+PyAttackForge is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+from .client import PyAttackForgeClient
+
+__all__ = ["PyAttackForgeClient"]

{pyattackforge-0.1.3 → pyattackforge-0.1.8}/pyattackforge/client.py

@@ -31,7 +31,7 @@ class PyAttackForgeClient:
     Supports dry-run mode for testing without making real API calls.
     """
 
-    def upsert_finding_for_project(
+    def upsert_finding_for_project(  # noqa: C901
         self,
         project_id: str,
         title: str,
@@ -87,20 +87,29 @@ class PyAttackForgeClient:
         asset_names = []
         for asset in affected_assets:
             name = asset["name"] if isinstance(asset, dict) and "name" in asset else asset
-            asset_obj = self.get_asset_by_name(name)
-            #if not asset_obj:
-            #    try:
-            #        asset_obj = self.create_asset({"name": name})
-            #    except Exception as e:
-            #        raise RuntimeError(f"Asset '{name}' does not exist and could not be created: {e}")
+            self.get_asset_by_name(name)
+            # if not asset_obj:
+            #     try:
+            #         asset_obj = self.create_asset({"name": name})
+            #     except Exception as e:
+            #         raise RuntimeError(f"Asset '{name}' does not exist and could not be created: {e}")
             asset_names.append(name)
 
         # Fetch all findings for the project
         findings = self.get_findings_for_project(project_id)
-        print(f"[DEBUG] get_findings_for_project({project_id}) returned {len(findings)} findings:")
+        logger.debug(
+            "Found %s findings for project %s",
+            len(findings),
+            project_id
+        )
         for f in findings:
-            print(f"  - id={f.get('vulnerability_id')}, title={f.get('vulnerability_title')}, steps_to_reproduce={f.get('vulnerability_steps_to_reproduce')}")
-            print(f"    FULL FINDING: {f}")
+            logger.debug(
+                "Finding id=%s title=%s steps=%s",
+                f.get("vulnerability_id"),
+                f.get("vulnerability_title"),
+                f.get("vulnerability_steps_to_reproduce"),
+            )
+            logger.debug("Finding payload: %s", f)
         match = None
         for f in findings:
             if f.get("vulnerability_title") == title:
@@ -189,6 +198,68 @@ class PyAttackForgeClient:
                 "result": result,
             }
 
+    def _list_project_findings(
+        self,
+        project_id: str,
+        params: Optional[Dict[str, Any]] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Internal helper to fetch findings for a project with optional query params.
+        """
+        if not project_id:
+            raise ValueError("Missing required field: project_id")
+        resp = self._request(
+            "get",
+            f"/api/ss/project/{project_id}/vulnerabilities",
+            params=params or {},
+        )
+        if resp.status_code != 200:
+            raise RuntimeError(f"Failed to fetch findings: {resp.text}")
+        data = resp.json()
+        if isinstance(data, dict) and "vulnerabilities" in data:
+            findings = data.get("vulnerabilities") or []
+        elif isinstance(data, list):
+            findings = data
+        else:
+            findings = []
+        return findings if isinstance(findings, list) else []
+
+    def get_findings(
+        self,
+        project_id: str,
+        page: int = 1,
+        limit: int = 100,
+        priority: Optional[str] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Backwards-compatible listing of findings with optional pagination.
+
+        Args:
+            project_id (str): The project ID.
+            page (int, optional): 1-based page number. Defaults to 1.
+            limit (int, optional): Page size. Defaults to 100.
+            priority (str, optional): Filter by priority.
+
+        Returns:
+            list: Page of finding/vulnerability dicts.
+        """
+        if page < 1:
+            raise ValueError("page must be >= 1")
+        if limit < 1:
+            raise ValueError("limit must be >= 1")
+        params: Dict[str, Any] = {
+            "skip": (page - 1) * limit,
+            "limit": limit,
+            "page": page,
+        }
+        if priority:
+            params["priority"] = priority
+        findings = self._list_project_findings(project_id, params=params)
+        if len(findings) > limit:
+            start = (page - 1) * limit
+            findings = findings[start:start + limit]
+        return findings
+
     def get_findings_for_project(self, project_id: str, priority: Optional[str] = None) -> list:
         """
         Fetch all findings/vulnerabilities for a given project.
@@ -200,20 +271,8 @@ class PyAttackForgeClient:
         Returns:
             list: List of finding/vulnerability dicts.
         """
-        params = {}
-        if priority:
-            params["priority"] = priority
-        resp = self._request("get", f"/api/ss/project/{project_id}/vulnerabilities", params=params)
-        if resp.status_code != 200:
-            raise RuntimeError(f"Failed to fetch findings: {resp.text}")
-        # The response may have a "vulnerabilities" key or be a list directly
-        data = resp.json()
-        if isinstance(data, dict) and "vulnerabilities" in data:
-            return data["vulnerabilities"]
-        elif isinstance(data, list):
-            return data
-        else:
-            return []
+        params = {"priority": priority} if priority else None
+        return self._list_project_findings(project_id, params=params)
 
     def get_vulnerability(self, vulnerability_id: str) -> Dict[str, Any]:
         """
@@ -235,6 +294,50 @@ class PyAttackForgeClient:
             return data["vulnerability"]
         return data
 
+    def update_finding(
+        self,
+        vulnerability_id: str,
+        project_id: Optional[str] = None,
+        affected_assets: Optional[list] = None,
+        notes: Optional[list] = None,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Update an existing finding/vulnerability with the provided fields.
+
+        Args:
+            vulnerability_id (str): The vulnerability ID to update.
+            project_id (str, optional): Project ID when required by the API.
+            affected_assets (list, optional): List of asset names or dicts with 'name'/'assetName'.
+            notes (list, optional): Notes payload to set.
+            **kwargs: Any additional fields accepted by the AttackForge API.
+
+        Returns:
+            dict: API response body.
+        """
+        if not vulnerability_id:
+            raise ValueError("Missing required field: vulnerability_id")
+        payload: Dict[str, Any] = {}
+        if project_id:
+            payload["project_id"] = project_id
+        if affected_assets is not None:
+            asset_names = [
+                a.get("assetName") if isinstance(a, dict) and "assetName" in a
+                else a.get("name") if isinstance(a, dict) and "name" in a
+                else a
+                for a in affected_assets
+            ]
+            payload["affected_assets"] = [{"assetName": n} for n in asset_names if n]
+        if notes is not None:
+            payload["notes"] = notes
+        for key, value in (kwargs or {}).items():
+            if value is not None:
+                payload[key] = value
+        resp = self._request("put", f"/api/ss/vulnerability/{vulnerability_id}", json_data=payload)
+        if resp.status_code not in (200, 201):
+            raise RuntimeError(f"Failed to update finding: {resp.text}")
+        return resp.json()
+
     def add_note_to_finding(
         self,
         vulnerability_id: str,
@@ -293,6 +396,88 @@ class PyAttackForgeClient:
             raise RuntimeError(f"Failed to add note: {resp.text}")
         return resp.json()
 
+    def link_vulnerability_to_testcases(
+        self,
+        vulnerability_id: str,
+        testcase_ids: List[str],
+        project_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Link a vulnerability to one or more testcases.
+
+        Args:
+            vulnerability_id (str): The vulnerability ID.
+            testcase_ids (list): List of testcase IDs to link.
+            project_id (str, optional): Project ID if required by the API.
+
+        Returns:
+            dict: API response.
+        """
+        if not vulnerability_id:
+            raise ValueError("Missing required field: vulnerability_id")
+        if not testcase_ids:
+            raise ValueError("testcase_ids must contain at least one ID")
+        payload: Dict[str, Any] = {
+            "linked_testcases": testcase_ids,
+        }
+        if project_id:
+            payload["project_id"] = project_id
+        resp = self._request(
+            "put",
+            f"/api/ss/vulnerability/{vulnerability_id}",
+            json_data=payload,
+        )
+        if resp.status_code not in (200, 201):
+            raise RuntimeError(f"Failed to link vulnerability to testcases: {resp.text}")
+        return resp.json()
+
+    def get_testcases(self, project_id: str) -> List[Dict[str, Any]]:
+        """
+        Retrieve testcases for a project.
+
+        Args:
+            project_id (str): Project ID.
+
+        Returns:
+            list: List of testcase dicts.
+        """
+        if not project_id:
+            raise ValueError("Missing required field: project_id")
+        resp = self._request("get", f"/api/ss/project/{project_id}/testcases")
+        if resp.status_code not in (200, 201):
+            raise RuntimeError(f"Failed to fetch testcases: {resp.text}")
+        data = resp.json()
+        if isinstance(data, dict) and "testcases" in data:
+            return data.get("testcases", [])
+        if isinstance(data, list):
+            return data
+        return []
+
+    def get_testcase(self, project_id: str, testcase_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve a single testcase by ID.
+
+        Args:
+            project_id (str): Project ID.
+            testcase_id (str): Testcase ID.
+
+        Returns:
+            dict or None: Testcase details if found, else None.
+        """
+        if not project_id:
+            raise ValueError("Missing required field: project_id")
+        if not testcase_id:
+            raise ValueError("Missing required field: testcase_id")
+        resp = self._request("get", f"/api/ss/project/{project_id}/testcase/{testcase_id}")
+        if resp.status_code == 404:
+            return None
+        if resp.status_code not in (200, 201):
+            raise RuntimeError(f"Failed to fetch testcase: {resp.text}")
+        data = resp.json()
+        if isinstance(data, dict) and "testcase" in data:
+            return data["testcase"]
+        return data if isinstance(data, dict) else None
+
     def upload_finding_evidence(self, vulnerability_id: str, file_path: str) -> Dict[str, Any]:
         """
         Upload evidence to a finding/vulnerability.
@@ -363,6 +548,47 @@ class PyAttackForgeClient:
             raise RuntimeError(f"Testcase evidence upload failed: {resp.text}")
         return resp.json()
 
+    def add_note_to_testcase(
+        self,
+        project_id: str,
+        testcase_id: str,
+        note: str,
+        status: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Create a testcase note via the dedicated note endpoint, optionally updating status via update_testcase.
+
+        Args:
+            project_id (str): Project ID.
+            testcase_id (str): Testcase ID.
+            note (str): Note text to set in the details field.
+            status (str, optional): Status to set (e.g., "Tested").
+
+        Returns:
+            dict: API response.
+        """
+        if not project_id:
+            raise ValueError("Missing required field: project_id")
+        if not testcase_id:
+            raise ValueError("Missing required field: testcase_id")
+        if not note:
+            raise ValueError("Missing required field: note")
+        endpoint = f"/api/ss/project/{project_id}/testcase/{testcase_id}/note"
+        payload: Dict[str, Any] = {"note": note, "note_type": "PLAINTEXT"}
+        resp = self._request("post", endpoint, json_data=payload)
+        if resp.status_code not in (200, 201):
+            raise RuntimeError(f"Failed to add testcase note: {resp.text}")
+        result = resp.json()
+
+        # Optionally update status using update_testcase
+        if status:
+            try:
+                self.update_testcase(project_id, testcase_id, {"status": status})
+            except Exception:
+                # If status update fails, still return note creation response
+                pass
+        return result
+
     def assign_findings_to_testcase(
         self,
         project_id: str,
@@ -429,6 +655,53 @@ class PyAttackForgeClient:
             raise RuntimeError(f"Failed to update testcase: {resp.text}")
         return resp.json()
 
+    def add_findings_to_testcase(
+        self,
+        project_id: str,
+        testcase_id: str,
+        vulnerability_ids: List[str],
+        additional_fields: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """
+        Fetch a testcase, merge existing linked vulnerabilities with the provided list, and update it.
+
+        Args:
+            project_id (str): The project ID.
+            testcase_id (str): The testcase ID.
+            vulnerability_ids (list): List of vulnerability IDs to add.
+            additional_fields (dict, optional): Extra fields to include (e.g., status).
+
+        Returns:
+            dict: API response from the update.
+        """
+        if not project_id:
+            raise ValueError("Missing required field: project_id")
+        if not testcase_id:
+            raise ValueError("Missing required field: testcase_id")
+        if not vulnerability_ids:
+            raise ValueError("vulnerability_ids must contain at least one ID")
+
+        testcases = self.get_testcases(project_id)
+        testcase = next((t for t in testcases if t.get("id") == testcase_id), None)
+        if not testcase:
+            raise RuntimeError(f"Testcase '{testcase_id}' not found in project '{project_id}'")
+
+        existing_raw = testcase.get("linked_vulnerabilities", []) or []
+        existing_ids: List[str] = []
+        for item in existing_raw:
+            if isinstance(item, dict) and item.get("id"):
+                existing_ids.append(item["id"])
+            elif isinstance(item, str):
+                existing_ids.append(item)
+
+        return self.assign_findings_to_testcase(
+            project_id=project_id,
+            testcase_id=testcase_id,
+            vulnerability_ids=vulnerability_ids,
+            existing_linked_vulnerabilities=existing_ids,
+            additional_fields=additional_fields,
+        )
+
     def __init__(self, api_key: str, base_url: str = "https://demo.attackforge.com", dry_run: bool = False):
         """
         Initialize the PyAttackForgeClient.
@@ -550,14 +823,14 @@ class PyAttackForgeClient:
 
     def create_asset(self, asset_data: Dict[str, Any]) -> Dict[str, Any]:
         pass
-        #resp = self._request("post", "/api/ss/library/asset", json_data=asset_data)
-        #if resp.status_code in (200, 201):
-        #    asset = resp.json()
-        #    self._asset_cache = None  # Invalidate cache
-        #    return asset
-        #if "Asset Already Exists" in resp.text:
-        #    return self.get_asset_by_name(asset_data["name"])
-        #raise RuntimeError(f"Asset creation failed: {resp.text}")
+        # resp = self._request("post", "/api/ss/library/asset", json_data=asset_data)
+        # if resp.status_code in (200, 201):
+        #     asset = resp.json()
+        #     self._asset_cache = None  # Invalidate cache
+        #     return asset
+        # if "Asset Already Exists" in resp.text:
+        #     return self.get_asset_by_name(asset_data["name"])
+        # raise RuntimeError(f"Asset creation failed: {resp.text}")
 
     def get_project_by_name(self, name: str) -> Optional[Dict[str, Any]]:
         params = {
@@ -653,6 +926,7 @@ class PyAttackForgeClient:
         writeup_id: str,
         priority: str,
         affected_assets: Optional[list] = None,
+        linked_testcases: Optional[list] = None,
         **kwargs
     ) -> Dict[str, Any]:
         """
@@ -663,6 +937,7 @@ class PyAttackForgeClient:
             writeup_id (str): The writeup/library ID.
             priority (str): The priority.
             affected_assets (list, optional): List of affected asset objects or names.
+            linked_testcases (list, optional): List of testcase IDs to link.
             **kwargs: Additional fields.
 
         Returns:
@@ -684,6 +959,8 @@ class PyAttackForgeClient:
                 for asset in affected_assets
             ]
             payload["affected_assets"] = [{"assetName": n} for n in asset_names]
+        if linked_testcases:
+            payload["linked_testcases"] = linked_testcases
         payload.update(kwargs)
         resp = self._request("post", "/api/ss/vulnerability-with-library", json_data=payload)
         if resp.status_code in (200, 201):
@@ -749,9 +1026,9 @@ class PyAttackForgeClient:
             name = asset["assetName"] if isinstance(asset, dict) and "assetName" in asset \
                 else asset["name"] if isinstance(asset, dict) and "name" in asset \
                 else asset
-            asset_obj = self.get_asset_by_name(name)
-            #if not asset_obj:
-            #    asset_obj = self.create_asset({"name": name})
+            self.get_asset_by_name(name)
+            # if not asset_obj:
+            #     asset_obj = self.create_asset({"name": name})
             asset_names.append(name)
         # Ensure all assets are in project scope
         scope = self.get_project_scope(project_id)
@@ -808,7 +1085,6 @@ class PyAttackForgeClient:
         )
         return result
 
-
     def create_vulnerability_old(
         self,
         project_id: str,
@@ -906,6 +1182,7 @@ class PyAttackForgeClient:
             return resp.json()
         raise RuntimeError(f"Vulnerability creation failed: {resp.text}")
 
+
 class DummyResponse:
     def __init__(self) -> None:
         self.status_code = 200

{pyattackforge-0.1.3 → pyattackforge-0.1.8}/pyattackforge.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyattackforge
-Version: 0.1.3
+Version: 0.1.8
 Summary: Python wrapper for the AttackForge API
 Home-page: https://github.com/Tantalum-Labs/PyAttackForge
 Author: Shane S
@@ -40,6 +40,7 @@ A lightweight Python library for interacting with the AttackForge API.
 - Create findings from existing writeups by passing a `writeup_id`
 - Upload evidence to findings or testcases
 - Update/assign testcases to link findings or add notes
+- Link vulnerabilities to testcases via the client
 - Dry-run mode for testing
 
 ---
@@ -112,14 +113,20 @@ client.create_vulnerability(
 
 ### Creating a finding from an existing writeup
 
-If you already have a writeup/library entry and just need to create a finding bound to it, you can either pass `writeup_id` to `create_vulnerability` (as above) or call `create_finding_from_writeup` directly:
+If you already have a writeup/library entry and just need to create a finding bound to it, you can either pass `writeup_id` to `create_vulnerability` (as above) or call `create_finding_from_writeup` directly. Prefer the 24-character writeup id (`id` / `_id`); if only a numeric `reference_id` is available, use that. You can also specify the library key (e.g., `approved_writeups`, `Main Vulnerabilities`):
 
 ```python
 client.create_finding_from_writeup(
     project_id="abc123",
-    writeup_id="68e92c7a821c05c8405a8003",
+    writeup_id="68e92c7a821c05c8405a8003",  # writeup id
+    library="approved_writeups",             # optional: library key/name
     priority="High",
-    affected_assets=[{"name": "ssh-prod-1"}]
+    affected_assets=[{"name": "ssh-prod-1"}],
+    linked_testcases=["5e8017d2e1385f0c58e8f4f8"],  # optional: link testcases at creation
+    likelihood_of_exploitation=5,
+    steps_to_reproduce="1. Do something\n2. Observe result",
+    notes=[{"note": "Created via API", "type": "PLAINTEXT"}],
+    tags=["automation"]
 )
 ```
 
@@ -152,21 +159,49 @@ client.add_note_to_finding(
 
 Add a note/update to a testcase (PUT to the testcase endpoint):
 ```python
-client.update_testcase(
+client.add_note_to_testcase(
     project_id="abc123",
     testcase_id="5e8017d2e1385f0c58e8f4f8",
-    update_fields={
-        "details": "Observed during retest on 2025-09-19."
-    }
+    note="Observed during retest on 2025-09-19.",
+    status="Tested"  # optional
 )
 ```
 
-Associate findings to a testcase (merges with existing linked vulnerabilities if provided):
+Associate findings to a testcase:
 ```python
 client.assign_findings_to_testcase(
     project_id="abc123",
     testcase_id="5e8017d2e1385f0c58e8f4f8",
-    vulnerability_ids=["66849b77950ab45e68fc7b48", "6768d29db1782d7362a2df5f"]
+    vulnerability_ids=["66849b77950ab45e68fc7b48", "6768d29db1782d7362a2df5f"],
+    additional_fields={"status": "Tested"} # optional
+)
+```
+Or link from the vulnerability side using its update endpoint:
+```python
+client.link_vulnerability_to_testcases(
+    vulnerability_id="69273ef0f4a7c85d03930667",
+    testcase_ids=["5e8017d2e1385f0c58e8f4f8"],
+    project_id="abc123",  # optional
+)
+```
+
+Fetch project testcases:
+```python
+testcases = client.get_testcases("abc123")
+```
+
+Fetch a single testcase (if supported in your tenant):
+```python
+testcase = client.get_testcase("abc123", "5e8017d2e1385f0c58e8f4f8")
+```
+
+Merge and add findings to a testcase in one call:
+```python
+client.add_findings_to_testcase(
+    project_id="abc123",
+    testcase_id="5e8017d2e1385f0c58e8f4f8",
+    vulnerability_ids=["69273ef0f4a7c85d03930667"],
+    additional_fields={"status": "Tested"} # optional
 )
 ```
 
@@ -212,32 +247,7 @@ See the source code for full details and docstrings.
 - `create_vulnerability(
       project_id: str,
       title: str,
-      affected_asset_name: str,
-      priority: str,
-      likelihood_of_exploitation: int,
-      description: str,
-      attack_scenario: str,
-      remediation_recommendation: str,
-      steps_to_reproduce: str,
-      tags: Optional[list] = None,
-      notes: Optional[list] = None,
-      is_zeroday: bool = False,
-      is_visible: bool = True,
-      import_to_library: Optional[str] = None,
-      import_source: Optional[str] = None,
-      import_source_id: Optional[str] = None,
-      custom_fields: Optional[list] = None,
-      linked_testcases: Optional[list] = None,
-      custom_tags: Optional[list] = None,
-  ) -> dict`
-
-See the source code for full details and docstrings.
-
----
-- `create_vulnerability(
-      project_id: str,
-      title: str,
-      affected_asset_name: str,
+      affected_assets: list,
       priority: str,
       likelihood_of_exploitation: int,
       description: str,
@@ -254,7 +264,21 @@ See the source code for full details and docstrings.
       custom_fields: Optional[list] = None,
       linked_testcases: Optional[list] = None,
       custom_tags: Optional[list] = None,
+      writeup_custom_fields: Optional[list] = None,
   ) -> dict`
+- `create_finding_from_writeup(project_id: str, writeup_id: str, priority: str, affected_assets: Optional[list] = None, linked_testcases: Optional[list] = None, **kwargs) -> dict`
+- `get_findings_for_project(project_id: str, priority: Optional[str] = None) -> list`
+- `upsert_finding_for_project(...)`
+- `get_vulnerability(vulnerability_id: str) -> dict`
+- `add_note_to_finding(vulnerability_id: str, note: Any, note_type: str = "PLAINTEXT") -> dict`
+- `upload_finding_evidence(vulnerability_id: str, file_path: str) -> dict`
+- `upload_testcase_evidence(project_id: str, testcase_id: str, file_path: str) -> dict`
+- `get_testcases(project_id: str) -> list`
+- `get_testcase(project_id: str, testcase_id: str) -> dict or None`
+- `link_vulnerability_to_testcases(vulnerability_id: str, testcase_ids: List[str], project_id: Optional[str] = None) -> dict`
+- `assign_findings_to_testcase(project_id: str, testcase_id: str, vulnerability_ids: List[str], existing_linked_vulnerabilities: Optional[List[str]] = None, additional_fields: Optional[Dict[str, Any]] = None) -> dict`
+- `add_findings_to_testcase(project_id: str, testcase_id: str, vulnerability_ids: List[str], additional_fields: Optional[Dict[str, Any]] = None) -> dict`
+- `add_note_to_testcase(project_id: str, testcase_id: str, note: str, status: Optional[str] = None) -> dict`
 
 See the source code for full details and docstrings.
 

{pyattackforge-0.1.3 → pyattackforge-0.1.8}/setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 setup(
     name="pyattackforge",
-    version="0.1.3",
+    version="0.1.8",
     packages=find_packages(),
     install_requires=[
         "requests>=2.20.0"

{pyattackforge-0.1.3 → pyattackforge-0.1.8}/tests/test_client.py

@@ -90,7 +90,7 @@ class TestPyAttackForgeClient(unittest.TestCase):
         # Ensure the client uses the provided writeup_id and does not attempt to create/search
         self.client.get_all_writeups = lambda force_refresh=False: []
         self.client.find_writeup_in_cache = lambda title, library="Main Vulnerabilities": None
-        captured = {}
+        captured = {"endpoints": []}
 
         def fake_create_from_writeup(**kwargs):
             captured.update(kwargs)
@@ -154,15 +154,42 @@ class TestPyAttackForgeClient(unittest.TestCase):
         )
         self.assertIsInstance(finding, dict)
 
-    def test_get_findings_for_project_dry_run(self):
-        findings = self.client.get_findings_for_project("dummy_project")
-        self.assertIsInstance(findings, list)
-
-    def test_upsert_finding_for_project_create(self):
-        # Simulate no existing findings (should create new)
-        self.client.get_findings_for_project = lambda project_id: []
-        # Patch get_all_writeups to return a matching writeup for this test
-        self.client.get_all_writeups = (
+    def test_get_findings_for_project_dry_run(self):
+        findings = self.client.get_findings_for_project("dummy_project")
+        self.assertIsInstance(findings, list)
+
+    def test_get_findings_with_pagination(self):
+        captured = {}
+
+        class Resp:
+            status_code = 200
+            text = "OK"
+
+            def json(self):
+                return {"vulnerabilities": [{"id": i} for i in range(10)]}
+
+        def fake_request(method, endpoint, json_data=None, params=None, files=None, data=None, headers_override=None):
+            captured["params"] = params
+            captured["endpoint"] = endpoint
+            return Resp()
+
+        self.client._request = fake_request
+        findings = self.client.get_findings("proj1", page=2, limit=3, priority="High")
+        self.assertEqual(len(findings), 3)
+        self.assertEqual(findings[0]["id"], 3)
+        self.assertEqual(captured["params"]["priority"], "High")
+        self.assertEqual(captured["params"]["skip"], 3)
+        self.assertEqual(captured["params"]["limit"], 3)
+        self.assertEqual(captured["params"]["page"], 2)
+        self.assertEqual(captured["endpoint"], "/api/ss/project/proj1/vulnerabilities")
+        with self.assertRaises(ValueError):
+            self.client.get_findings("proj1", page=0)
+
+    def test_upsert_finding_for_project_create(self):
+        # Simulate no existing findings (should create new)
+        self.client.get_findings_for_project = lambda project_id: []
+        # Patch get_all_writeups to return a matching writeup for this test
+        self.client.get_all_writeups = (
             lambda force_refresh=False: [
                 {
                     "title": "UnitTest Finding",
@@ -203,24 +230,26 @@ class TestPyAttackForgeClient(unittest.TestCase):
         }
         self.client.get_findings_for_project = lambda project_id: [existing_finding]
         # Patch get_all_writeups to return a matching writeup for this test
-        self.client.get_all_writeups = (
-            lambda force_refresh=False: [
-                {
-                    "title": "UnitTest Finding",
-                    "belongs_to_library": "Main Vulnerabilities",
-                    "reference_id": "dummy_writeup_id",
-                }
-            ]
-        )
-        self.client.create_writeup = lambda **kwargs: {"reference_id": "dummy_writeup_id"}
-        # Patch _request to simulate API update response
-        class Resp:
-            status_code = 200
-            def json(self):
-                return {"updated": True}
-            text = "OK"
-        self.client._request = (
-            lambda method, endpoint, json_data=None, params=None: Resp()
+        self.client.get_all_writeups = (
+            lambda force_refresh=False: [
+                {
+                    "title": "UnitTest Finding",
+                    "belongs_to_library": "Main Vulnerabilities",
+                    "reference_id": "dummy_writeup_id",
+                }
+            ]
+        )
+        self.client.create_writeup = lambda **kwargs: {"reference_id": "dummy_writeup_id"}
+        # Patch _request to simulate API update response
+
+        class Resp:
+            status_code = 200
+
+            def json(self):
+                return {"updated": True}
+            text = "OK"
+        self.client._request = (
+            lambda method, endpoint, json_data=None, params=None: Resp()
         )
         result = self.client.upsert_finding_for_project(
             project_id="dummy_project",
@@ -339,7 +368,7 @@ class TestPyAttackForgeClient(unittest.TestCase):
             os.remove(evidence_path)
 
     def test_assign_findings_to_testcase_merges(self):
-        captured = {}
+        captured = {"endpoints": []}
 
         def fake_update(project_id, testcase_id, update_fields):
             captured["payload"] = update_fields
@@ -360,11 +389,12 @@ class TestPyAttackForgeClient(unittest.TestCase):
         self.client.get_vulnerability = lambda vid: {
             "vulnerability_notes": [{"note": "Existing note", "type": "PLAINTEXT"}]
         }
-        captured = {}
+        captured = {"endpoints": []}
 
         class Resp:
             status_code = 200
             text = "OK"
+
             def json(self):
                 return {"ok": True}
 
@@ -376,7 +406,148 @@ class TestPyAttackForgeClient(unittest.TestCase):
         self.assertIsInstance(result, dict)
         notes = captured["json_data"].get("notes", [])
         self.assertEqual(len(notes), 1)
-
+
+    def test_update_finding(self):
+        captured = {}
+
+        class Resp:
+            status_code = 200
+            text = "OK"
+
+            def json(self):
+                return {"updated": True}
+
+        def fake_request(method, endpoint, json_data=None, params=None, files=None, data=None, headers_override=None):
+            captured["endpoint"] = endpoint
+            captured["json_data"] = json_data
+            return Resp()
+
+        self.client._request = fake_request
+        resp = self.client.update_finding(
+            vulnerability_id="v-1",
+            project_id="p-1",
+            affected_assets=[{"name": "asset-a"}, {"assetName": "asset-b"}, "asset-c"],
+            notes=[{"note": "n1", "type": "PLAINTEXT"}],
+            custom="field",
+        )
+        self.assertIsInstance(resp, dict)
+        payload = captured["json_data"]
+        self.assertEqual(payload["project_id"], "p-1")
+        self.assertEqual(
+            payload["affected_assets"],
+            [{"assetName": "asset-a"}, {"assetName": "asset-b"}, {"assetName": "asset-c"}],
+        )
+        self.assertEqual(payload["notes"], [{"note": "n1", "type": "PLAINTEXT"}])
+        self.assertEqual(payload["custom"], "field")
+        self.assertEqual(captured["endpoint"], "/api/ss/vulnerability/v-1")
+        with self.assertRaises(ValueError):
+            self.client.update_finding("", project_id="p-1")
+
+    def test_get_testcases(self):
+        # DummyResponse returns {}, so should yield an empty list without raising
+        cases = self.client.get_testcases("proj1")
+        self.assertIsInstance(cases, list)
+        self.assertEqual(cases, [])
+
+    def test_get_testcase(self):
+        class Resp:
+            def __init__(self, status_code, body):
+                self.status_code = status_code
+                self._body = body
+
+            def json(self):
+                return self._body
+            text = "resp"
+
+        calls = []
+
+        def fake_request(method, endpoint, json_data=None, params=None, files=None, data=None, headers_override=None):
+            calls.append(endpoint)
+            if "tc-ok" in endpoint:
+                return Resp(200, {"testcase": {"id": "tc-ok", "status": "Not Tested"}})
+            return Resp(404, {})
+
+        self.client._request = fake_request
+        tc_none = self.client.get_testcase("proj", "tc-missing")
+        self.assertIsNone(tc_none)
+        tc = self.client.get_testcase("proj", "tc-ok")
+        self.assertIsInstance(tc, dict)
+        self.assertEqual(tc.get("id"), "tc-ok")
+
+    def test_add_note_to_testcase(self):
+        captured = {"endpoints": []}
+
+        class Resp:
+            status_code = 200
+
+            def json(self):
+                return {"status": "Testcase Note Created"}
+
+        def fake_request(method, endpoint, json_data=None, params=None, files=None, data=None, headers_override=None):
+            captured["endpoints"].append((endpoint, json_data))
+            return Resp()
+
+        self.client._request = fake_request
+        resp = self.client.add_note_to_testcase("proj1", "tc1", "Note text", status="Tested")
+        self.assertIsInstance(resp, dict)
+        note_calls = [c for c in captured["endpoints"] if "/note" in c[0]]
+        self.assertTrue(note_calls)
+        self.assertEqual(note_calls[0][1]["note"], "Note text")
+        self.assertEqual(note_calls[0][1]["note_type"], "PLAINTEXT")
+
+    def test_link_vulnerability_to_testcases(self):
+        captured = {}
+
+        class Resp:
+            status_code = 200
+            text = "OK"
+
+            def json(self):
+                return {"linked": True}
+
+        def fake_request(method, endpoint, json_data=None, params=None, files=None, data=None, headers_override=None):
+            captured["method"] = method
+            captured["endpoint"] = endpoint
+            captured["json_data"] = json_data
+            return Resp()
+
+        self.client._request = fake_request
+        resp = self.client.link_vulnerability_to_testcases("v1", ["tc1", "tc2"], project_id="proj1")
+        self.assertIsInstance(resp, dict)
+        self.assertEqual(captured["endpoint"], "/api/ss/vulnerability/v1")
+        self.assertEqual(captured["json_data"]["linked_testcases"], ["tc1", "tc2"])
+        self.assertEqual(captured["json_data"]["project_id"], "proj1")
+
+    def test_add_findings_to_testcase(self):
+        captured = {}
+        # Simulate existing testcase with one linked vuln (as dict)
+        self.client.get_testcases = lambda project_id: [
+            {
+                "id": "tc1",
+                "linked_vulnerabilities": [{"id": "existing"}],
+            }
+        ]
+
+        def fake_assign(project_id, testcase_id, vulnerability_ids, existing_linked_vulnerabilities=None, additional_fields=None):
+            captured["project_id"] = project_id
+            captured["testcase_id"] = testcase_id
+            captured["vulnerability_ids"] = vulnerability_ids
+            captured["existing_linked_vulnerabilities"] = existing_linked_vulnerabilities
+            captured["additional_fields"] = additional_fields
+            return {"assigned": True}
+
+        self.client.assign_findings_to_testcase = fake_assign
+        resp = self.client.add_findings_to_testcase(
+            "proj1",
+            "tc1",
+            ["new1", "new2"],
+            additional_fields={"status": "Tested"},
+        )
+        self.assertIsInstance(resp, dict)
+        self.assertEqual(captured["existing_linked_vulnerabilities"], ["existing"])
+        self.assertEqual(captured["vulnerability_ids"], ["new1", "new2"])
+        self.assertEqual(captured["additional_fields"], {"status": "Tested"})
+
 
 if __name__ == "__main__":
     unittest.main()