pyattackforge 0.1.3__py3-none-any.whl → 0.1.7__py3-none-any.whl

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/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:
@@ -293,6 +302,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 +454,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 +561,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 +729,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 +832,7 @@ class PyAttackForgeClient:
         writeup_id: str,
         priority: str,
         affected_assets: Optional[list] = None,
+        linked_testcases: Optional[list] = None,
         **kwargs
     ) -> Dict[str, Any]:
         """
@@ -663,6 +843,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 +865,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 +932,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 +991,6 @@ class PyAttackForgeClient:
         )
         return result
 
-
     def create_vulnerability_old(
         self,
         project_id: str,
@@ -906,6 +1088,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.dist-info → pyattackforge-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyattackforge
-Version: 0.1.3
+Version: 0.1.7
 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.7.dist-info/RECORD

@@ -0,0 +1,8 @@
+pyattackforge/__init__.py,sha256=r4FT68O03ZuyQJ70RLwlLDqp7AqNTTcMm-FLU4zHwIU,807
+pyattackforge/client.py,sha256=Z-rGVPYpmG8x4N79Izv1Q4sC3-aX40Lf7XdT8kidor8,45464
+pyattackforge/prev_client.py,sha256=Wu-5BVCJbLkWBsxjQmL3nZoyWqAqOACKM9ScXfl6C5o,14277
+pyattackforge-0.1.7.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+pyattackforge-0.1.7.dist-info/METADATA,sha256=fU2BIgYD4OqcLMIqeLWdtvBfHBf09M_OnAtUkkhnbHc,10754
+pyattackforge-0.1.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+pyattackforge-0.1.7.dist-info/top_level.txt,sha256=1rDeMkWvFWuX3MS8V65no7KuybYyvtfIgbYSt5m_uPU,14
+pyattackforge-0.1.7.dist-info/RECORD,,

pyattackforge-0.1.3.dist-info/RECORD

@@ -1,8 +0,0 @@
-pyattackforge/__init__.py,sha256=xZxubdjv_Fc1bRqfVRMR4j9SCTLXg9TfckWn9CQHH_E,839
-pyattackforge/client.py,sha256=gPd5CwP714NKe0IwhmKUrCMklQg1ARVOfcB4ibAne7o,38753
-pyattackforge/prev_client.py,sha256=Wu-5BVCJbLkWBsxjQmL3nZoyWqAqOACKM9ScXfl6C5o,14277
-pyattackforge-0.1.3.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-pyattackforge-0.1.3.dist-info/METADATA,sha256=xzIKr5FGJ37Enc5QaCPyiVj0wFlgyUdR63Wq_u_M9sE,8770
-pyattackforge-0.1.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-pyattackforge-0.1.3.dist-info/top_level.txt,sha256=1rDeMkWvFWuX3MS8V65no7KuybYyvtfIgbYSt5m_uPU,14
-pyattackforge-0.1.3.dist-info/RECORD,,