pyattackforge 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

pyattackforge/client.py

@@ -11,6 +11,7 @@ from .transport import AttackForgeTransport
 from .resources import (
     AssetsResource,
     ProjectsResource,
+    GroupsResource,
     FindingsResource,
     WriteupsResource,
     TestcasesResource,
@@ -51,6 +52,7 @@ class AttackForgeClient:
 
         self.assets = AssetsResource(self._transport)
         self.projects = ProjectsResource(self._transport)
+        self.groups = GroupsResource(self._transport)
         self.findings = FindingsResource(self._transport)
         self.writeups = WriteupsResource(self._transport)
         self.testcases = TestcasesResource(self._transport)

pyattackforge/config.py

@@ -16,7 +16,7 @@ class ClientConfig:
     timeout: float = 30.0
     max_retries: int = 3
     backoff_factor: float = 0.5
-    user_agent: str = "pyattackforge/0.0.1"
+    user_agent: str = "pyattackforge/0.2.2"
     http2: bool = True
     # Default visibility for newly created findings. False = pending/hidden.
     default_findings_visible: bool = False

pyattackforge/resources/__init__.py

@@ -2,6 +2,7 @@
 
 from .assets import AssetsResource
 from .projects import ProjectsResource
+from .groups import GroupsResource
 from .findings import FindingsResource
 from .writeups import WriteupsResource
 from .testcases import TestcasesResource
@@ -13,6 +14,7 @@ from .reports import ReportsResource
 __all__ = [
     "AssetsResource",
     "ProjectsResource",
+    "GroupsResource",
     "FindingsResource",
     "WriteupsResource",
     "TestcasesResource",

pyattackforge/resources/findings.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 from typing import Any, Dict, Iterable, List, Optional, Sequence
 import os
 from datetime import datetime
+import time
 
 from ..cache import TTLCache
 from .base import BaseResource
@@ -247,18 +248,29 @@ class FindingsResource(BaseResource):
             return {"action": "noop", "existing": match}
 
         merged = sorted(existing_assets.union(new_assets))
-        asset_ids = self._resolve_project_asset_ids(project_id, merged)
-        payload_assets = []
-        for name in merged:
-            asset_id = asset_ids.get(name)
-            if asset_id:
-                payload_assets.append({"assetId": asset_id})
-            else:
-                payload_assets.append({"assetName": name})
-        payload = {"affected_assets": payload_assets}
+        asset_ids = {}
+        asset_ids.update(self._extract_asset_id_map_from_finding(match))
+        asset_ids.update(self._map_names_to_ids_from_payload(affected_assets, create_payload))
         if update_payload:
-            payload.update(update_payload)
-        result = self.update_vulnerability(match.get("vulnerability_id") or match.get("id"), payload)
+            asset_ids.update(self._map_names_to_ids_from_payload(affected_assets, update_payload))
+        resolved = self._resolve_project_asset_ids_with_retry(project_id, merged)
+        asset_ids.update(resolved)
+        payload_assets = self._build_affected_assets_payload(merged, asset_ids)
+        payload = dict(update_payload or {})
+        payload.pop("affected_assets", None)
+        payload.pop("vulnerability_affected_assets", None)
+        payload.pop("projectId", None)
+        payload.pop("project_id", None)
+        payload["affected_assets"] = payload_assets
+        vuln_id = match.get("vulnerability_id") or match.get("id")
+        result = self.update_vulnerability(vuln_id, payload)
+        if missing:
+            result = self._ensure_assets_on_vulnerability(
+                vulnerability_id=vuln_id,
+                project_id=project_id,
+                desired_assets=merged,
+                asset_ids=asset_ids,
+            )
         return {"action": "update", "result": result, "added_assets": missing}
 
     def _normalize_title(self, title: str) -> str:
@@ -376,11 +388,143 @@ class FindingsResource(BaseResource):
             if not isinstance(entry, dict):
                 continue
             name = entry.get("name")
-            asset_id = entry.get("asset_id") or entry.get("assetId")
+            asset_id = entry.get("id") or entry.get("asset_id") or entry.get("assetId")
             if isinstance(name, str) and isinstance(asset_id, str):
+                stripped = name.strip()
+                if stripped:
+                    mapping[stripped] = asset_id
+        return mapping
+
+    def _resolve_project_asset_ids_with_retry(
+        self,
+        project_id: str,
+        asset_names: Sequence[str],
+        *,
+        attempts: int = 3,
+        delay: float = 0.5,
+    ) -> Dict[str, str]:
+        seen: set = set()
+        unique_names: List[str] = []
+        for value in asset_names:
+            if not isinstance(value, str):
+                continue
+            name = value.strip()
+            if not name or name in seen:
+                continue
+            seen.add(name)
+            unique_names.append(name)
+        if not unique_names:
+            return {}
+        mapping = self._resolve_project_asset_ids(project_id, unique_names)
+        missing = [name for name in unique_names if name not in mapping]
+        if not missing or attempts <= 1:
+            return mapping
+        for _ in range(attempts - 1):
+            time.sleep(delay)
+            refreshed = self._resolve_project_asset_ids(project_id, missing)
+            mapping.update(refreshed)
+            missing = [name for name in missing if name not in mapping]
+            if not missing:
+                break
+        return mapping
+
+    def _build_affected_assets_payload(self, asset_names: Sequence[str], asset_ids: Dict[str, str]) -> List[Dict[str, Any]]:
+        payload_assets: List[Dict[str, Any]] = []
+        for name in asset_names:
+            if not isinstance(name, str):
+                continue
+            stripped = name.strip()
+            if not stripped:
+                continue
+            asset_id = asset_ids.get(stripped)
+            if asset_id:
+                payload_assets.append({"assetId": asset_id})
+            else:
+                payload_assets.append({"assetName": stripped})
+        return payload_assets
+
+    def _ensure_assets_on_vulnerability(
+        self,
+        *,
+        vulnerability_id: str,
+        project_id: str,
+        desired_assets: Sequence[str],
+        asset_ids: Dict[str, str],
+        attempts: int = 3,
+        delay: float = 1.0,
+    ) -> Any:
+        last_result: Any = None
+        for _ in range(max(attempts, 1)):
+            try:
+                vuln = self.get_vulnerability(vulnerability_id)
+            except Exception:
+                break
+            current_assets = self._extract_asset_names_from_finding(vuln)
+            missing = [name for name in desired_assets if name not in current_assets]
+            if not missing:
+                return last_result if last_result is not None else vuln
+            refreshed = self._resolve_project_asset_ids_with_retry(
+                project_id, missing, attempts=2, delay=delay
+            )
+            asset_ids.update(refreshed)
+            payload_assets = self._build_affected_assets_payload(desired_assets, asset_ids)
+            payload = {"affected_assets": payload_assets}
+            last_result = self.update_vulnerability(vulnerability_id, payload)
+            time.sleep(delay)
+        return last_result
+
+    def _extract_asset_id_map_from_finding(self, finding: Dict[str, Any]) -> Dict[str, str]:
+        raw = finding.get("vulnerability_affected_assets") or finding.get("affected_assets") or []
+        return self._extract_asset_id_map(raw if isinstance(raw, list) else [])
+
+    def _extract_asset_id_map(self, assets: Sequence[Any]) -> Dict[str, str]:
+        mapping: Dict[str, str] = {}
+        for item in assets:
+            if not isinstance(item, dict):
+                continue
+            name = None
+            asset_id = None
+            for key in ("assetName", "name"):
+                value = item.get(key)
+                if isinstance(value, str) and value.strip():
+                    name = value.strip()
+                    break
+            asset_obj = item.get("asset")
+            if isinstance(asset_obj, dict):
+                obj_name = asset_obj.get("name")
+                if isinstance(obj_name, str) and obj_name.strip():
+                    name = obj_name.strip()
+                obj_id = asset_obj.get("id") or asset_obj.get("asset_id") or asset_obj.get("assetId")
+                if isinstance(obj_id, str) and obj_id.strip():
+                    asset_id = obj_id.strip()
+            for key in ("assetId", "asset_id", "id"):
+                value = item.get(key)
+                if isinstance(value, str) and value.strip():
+                    asset_id = value.strip()
+                    break
+            if name and asset_id:
                 mapping[name] = asset_id
         return mapping
 
+    def _map_names_to_ids_from_payload(
+        self, asset_names: Sequence[Any], payload: Optional[Dict[str, Any]]
+    ) -> Dict[str, str]:
+        if not isinstance(payload, dict):
+            return {}
+        payload_assets = payload.get("affected_assets") or payload.get("vulnerability_affected_assets") or []
+        if not isinstance(payload_assets, list):
+            return {}
+        mapping = self._extract_asset_id_map(payload_assets)
+        names_list = [name for name in asset_names if isinstance(name, str) and name.strip()]
+        if names_list and len(names_list) == len(payload_assets):
+            for name, entry in zip(names_list, payload_assets):
+                if name in mapping or not isinstance(entry, dict):
+                    continue
+                asset_id = entry.get("assetId") or entry.get("asset_id") or entry.get("id")
+                if isinstance(asset_id, str) and asset_id.strip():
+                    mapping[name] = asset_id.strip()
+        return mapping
+
     def _enforce_vulnerability_evidence_fifo(
         self, vulnerability_id: str, keep_last: int, *, project_id: Optional[str] = None
     ) -> List[str]:

pyattackforge/resources/groups.py

@@ -0,0 +1,20 @@
+"""Resource: groups."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from .base import BaseResource
+
+
+class GroupsResource(BaseResource):
+    """Groups API resource wrapper."""
+
+    def create_group(self, payload: Dict[str, Any]) -> Any:
+        return self._post("/api/ss/group", json=payload)
+
+    def get_groups(self, params: Optional[Dict[str, Any]] = None) -> Any:
+        return self._get("/api/ss/groups", params=params)
+
+    def get_group(self, group_id: str, params: Optional[Dict[str, Any]] = None) -> Any:
+        return self._get(f"/api/ss/group/{group_id}", params=params)

pyattackforge/resources/projects.py

@@ -130,6 +130,46 @@ class ProjectsResource(BaseResource):
     def update_user_access_on_project(self, project_id: str, user_id: str, payload: Dict[str, Any]) -> Any:
         return self._put(f"/api/ss/project/{project_id}/access/{user_id}", json=payload)
 
+    def add_project_to_group(
+        self,
+        project_id: str,
+        group_id: str,
+        *,
+        project_data: Optional[Dict[str, Any]] = None,
+    ) -> Any:
+        """
+        Convenience helper to add a project to a group by updating the project's groups list.
+        """
+        if not group_id:
+            raise ValueError("group_id is required")
+        data = project_data if project_data is not None else self.get_project(project_id)
+        project = data
+        if isinstance(project, dict) and isinstance(project.get("data"), dict):
+            project = project["data"]
+        if isinstance(project, dict) and isinstance(project.get("project"), dict):
+            project = project["project"]
+        groups = project.get("groups") or project.get("project_groups") or []
+        group_ids: List[str] = []
+        if isinstance(groups, list):
+            for entry in groups:
+                if isinstance(entry, str) and entry:
+                    group_ids.append(entry)
+                    continue
+                if isinstance(entry, dict):
+                    candidate = entry.get("id") or entry.get("group_id") or entry.get("_id")
+                    if isinstance(candidate, str) and candidate:
+                        group_ids.append(candidate)
+        seen: set = set()
+        deduped: List[str] = []
+        for value in group_ids:
+            if value in seen:
+                continue
+            seen.add(value)
+            deduped.append(value)
+        if group_id not in seen:
+            deduped.append(group_id)
+        return self.update_project(project_id, {"groups": deduped})
+
     def extract_projects_list(self, projects_data: Any) -> List[Dict[str, Any]]:
         if not isinstance(projects_data, dict):
             return []

pyattackforge/transport.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Iterator, Tuple
 import time
 import httpx
 
@@ -52,6 +52,42 @@ class AttackForgeTransport:
             )
         return self._http2_client
 
+    def _iter_file_objects(self, files: Dict[str, Any]) -> Iterator[Any]:
+        for value in files.values():
+            file_obj = None
+            if hasattr(value, "read"):
+                file_obj = value
+            elif isinstance(value, tuple) and len(value) >= 2 and hasattr(value[1], "read"):
+                file_obj = value[1]
+            if file_obj is not None:
+                yield file_obj
+
+    def _capture_file_positions(self, files: Dict[str, Any]) -> Tuple[Dict[int, Optional[int]], bool]:
+        positions: Dict[int, Optional[int]] = {}
+        reusable = True
+        for file_obj in self._iter_file_objects(files):
+            try:
+                if hasattr(file_obj, "seekable") and not file_obj.seekable():
+                    reusable = False
+                positions[id(file_obj)] = file_obj.tell()
+            except Exception:
+                positions[id(file_obj)] = None
+                reusable = False
+        return positions, reusable
+
+    def _rewind_files(self, files: Dict[str, Any], positions: Dict[int, Optional[int]]) -> bool:
+        ok = True
+        for file_obj in self._iter_file_objects(files):
+            pos = positions.get(id(file_obj))
+            if pos is None:
+                ok = False
+                continue
+            try:
+                file_obj.seek(pos)
+            except Exception:
+                ok = False
+        return ok
+
     def request(
         self,
         method: str,
@@ -68,8 +104,16 @@ class AttackForgeTransport:
         attempt = 0
         max_retries = self._config.max_retries if retries is None else retries
         backoff = self._config.backoff_factor
+        file_positions: Optional[Dict[int, Optional[int]]] = None
+        files_reusable = True
+        if files:
+            file_positions, files_reusable = self._capture_file_positions(files)
+            if not files_reusable:
+                max_retries = 0
         while True:
             try:
+                if files and attempt > 0 and files_reusable:
+                    self._rewind_files(files, file_positions or {})
                 request_headers = None
                 if headers:
                     request_headers = headers
@@ -96,19 +140,21 @@ class AttackForgeTransport:
                 continue
 
             if response.status_code >= 500 and files:
-                try:
-                    response = self._get_http2_client().request(
-                        method=method,
-                        url=path,
-                        params=params,
-                        json=json,
-                        files=files,
-                        data=data,
-                        headers=request_headers,
-                        timeout=timeout,
-                    )
-                except httpx.HTTPError:
-                    pass
+                if files_reusable:
+                    self._rewind_files(files, file_positions or {})
+                    try:
+                        response = self._get_http2_client().request(
+                            method=method,
+                            url=path,
+                            params=params,
+                            json=json,
+                            files=files,
+                            data=data,
+                            headers=request_headers,
+                            timeout=timeout,
+                        )
+                    except httpx.HTTPError:
+                        pass
 
             if response.status_code in (429, 500, 502, 503, 504) and attempt < max_retries:
                 time.sleep(backoff * (2 ** attempt))