pyattackforge 0.0.1__py3-none-any.whl

pyattackforge/resources/notes.py

@@ -0,0 +1,139 @@
+"""Resource: notes."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+import os
+import base64
+
+from .base import BaseResource
+from ..exceptions import APIError
+
+
+class NotesResource(BaseResource):
+    """Notes API resource wrapper."""
+
+    def create_remediation_note(self, vulnerability_id: str, payload: Dict[str, Any]) -> Any:
+        return self._post(f"/api/ss/vulnerability/{vulnerability_id}/remediationNote", json=payload)
+
+    def update_remediation_note(self, vulnerability_id: str, remediation_note_id: str, payload: Dict[str, Any]) -> Any:
+        return self._put(
+            f"/api/ss/vulnerability/{vulnerability_id}/remediationNote/{remediation_note_id}",
+            json=payload,
+        )
+
+    def upload_remediation_note_file(self, vulnerability_id: str, remediation_note_id: str, file_path: str) -> Any:
+        if not os.path.isfile(file_path):
+            raise FileNotFoundError(file_path)
+        with open(file_path, "rb") as handle:
+            return self._post_files(
+                f"/api/ss/vulnerability/{vulnerability_id}/remediationNote/{remediation_note_id}/file",
+                files={"file": (os.path.basename(file_path), handle)},
+            )
+
+    def download_remediation_note_file(
+        self,
+        vulnerability_id: str,
+        remediation_note_id: str,
+        file_name: str,
+        *,
+        project_id: Optional[str] = None,
+        allow_report_fallback: bool = True,
+    ) -> Any:
+        """
+        Download a remediation note file.
+
+        If the SSAPI download endpoint fails, optionally fall back to report data
+        (base64) when available.
+        """
+        try:
+            return self._get(
+                f"/api/ss/vulnerability/{vulnerability_id}/remediationNote/{remediation_note_id}/file/{file_name}"
+            )
+        except APIError as exc:
+            if not allow_report_fallback:
+                raise
+            resolved_project_id = project_id or self._find_project_id_for_vulnerability(vulnerability_id)
+            if not resolved_project_id:
+                raise
+            report = self._post(
+                f"/api/ss/project/{resolved_project_id}/report/raw",
+                json={"excludeBinaries": False, "vulnerabilityIds": [vulnerability_id]},
+            )
+            file_entry = self._find_remediation_note_file_from_report(
+                report, vulnerability_id, remediation_note_id, file_name
+            )
+            if file_entry is None:
+                report = self._post(
+                    f"/api/ss/project/{resolved_project_id}/report/raw",
+                    json={"excludeBinaries": False},
+                )
+                file_entry = self._find_remediation_note_file_from_report(
+                    report, vulnerability_id, remediation_note_id, file_name
+                )
+            if isinstance(file_entry, dict):
+                payload = file_entry.get("fileBase64")
+                if isinstance(payload, str) and payload:
+                    try:
+                        padded = payload + ("=" * (-len(payload) % 4))
+                        return base64.b64decode(padded)
+                    except (ValueError, TypeError):
+                        pass
+            raise exc
+
+    def _find_project_id_for_vulnerability(self, vulnerability_id: str) -> Optional[str]:
+        data = self._get("/api/ss/projects-and-vulnerabilities")
+        projects = data.get("projects") if isinstance(data, dict) else None
+        if not isinstance(projects, list):
+            return None
+        for project in projects:
+            if not isinstance(project, dict):
+                continue
+            vulns = project.get("project_vulnerabilities")
+            if not isinstance(vulns, list):
+                continue
+            for vuln in vulns:
+                if not isinstance(vuln, dict):
+                    continue
+                vid = vuln.get("vulnerability_id") or vuln.get("id")
+                if vid == vulnerability_id:
+                    return project.get("project_id") or project.get("id")
+        return None
+
+    def _find_remediation_note_file_from_report(
+        self,
+        report: Any,
+        vulnerability_id: str,
+        remediation_note_id: str,
+        file_name: Optional[str],
+    ) -> Optional[Dict[str, Any]]:
+        if not isinstance(report, dict):
+            return None
+        vulnerabilities = report.get("vulnerabilities")
+        if not isinstance(vulnerabilities, list):
+            return None
+        for vuln in vulnerabilities:
+            if not isinstance(vuln, dict):
+                continue
+            vid = vuln.get("id") or vuln.get("vulnerability_id")
+            if vid != vulnerability_id:
+                continue
+            for asset in vuln.get("affected_assets") or []:
+                if not isinstance(asset, dict):
+                    continue
+                for note in asset.get("remediation_notes") or []:
+                    if not isinstance(note, dict):
+                        continue
+                    nid = note.get("id") or note.get("remediation_note_id") or note.get("note_id")
+                    if nid != remediation_note_id:
+                        continue
+                    files = note.get("remediation_note_files") or []
+                    if not isinstance(files, list):
+                        continue
+                    for entry in files:
+                        if not isinstance(entry, dict):
+                            continue
+                        entry_name = entry.get("fileName") or entry.get("file_name") or entry.get("name")
+                        if file_name is None or entry_name == file_name:
+                            return entry
+        return None

pyattackforge/resources/projects.py

@@ -0,0 +1,154 @@
+"""Resource: projects."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Iterable, List
+import os
+
+from ..cache import TTLCache
+from .base import BaseResource
+
+
+class ProjectsResource(BaseResource):
+    """Projects API resource wrapper."""
+
+    def __init__(self, transport) -> None:  # type: ignore[override]
+        super().__init__(transport)
+        self._cache = TTLCache(default_ttl=300.0)
+
+    def create_project(self, payload: Dict[str, Any]) -> Any:
+        return self._post("/api/ss/project", json=payload)
+
+    def get_project(self, project_id: str, params: Optional[Dict[str, Any]] = None, *, force_refresh: bool = False) -> Any:
+        cache_key = f"project:{project_id}"
+        if params is None and not force_refresh:
+            cached = self._cache.get(cache_key)
+            if cached is not None:
+                return cached
+        data = self._get(f"/api/ss/project/{project_id}", params=params)
+        if params is None:
+            self._cache.set(cache_key, data)
+        return data
+
+    def get_projects(self, params: Optional[Dict[str, Any]] = None) -> Any:
+        return self._get("/api/ss/projects", params=params)
+
+    def get_projects_and_vulnerabilities(self, params: Optional[Dict[str, Any]] = None) -> Any:
+        return self._get("/api/ss/projects-and-vulnerabilities", params=params)
+
+    def update_project(self, project_id: str, payload: Dict[str, Any]) -> Any:
+        data = self._put(f"/api/ss/project/{project_id}", json=payload)
+        self._cache.set(f"project:{project_id}", data)
+        return data
+
+    def archive_project(self, project_id: str) -> Any:
+        return self._put(f"/api/ss/project/{project_id}/archive")
+
+    def restore_project(self, project_id: str) -> Any:
+        return self._put(f"/api/ss/project/{project_id}/restore")
+
+    def destroy_projects(self, project_ids: Iterable[str], keep_logs: Optional[bool] = None) -> Any:
+        payload: Dict[str, Any] = {"project_ids": list(project_ids)}
+        if keep_logs is not None:
+            payload["keep_logs"] = keep_logs
+        return self._delete("/api/ss/project/destroy", json=payload)
+
+    def clone_project(self, project_id: str, payload: Optional[Dict[str, Any]] = None) -> Any:
+        return self._post(f"/api/ss/project/{project_id}/clone", json=payload or {})
+
+    def create_scope(self, project_id: str, payload: Dict[str, Any]) -> Any:
+        return self._post(f"/api/ss/project/{project_id}/assets", json=payload)
+
+    def update_scope(self, project_id: str, asset_id: str, payload: Dict[str, Any]) -> Any:
+        return self._put(f"/api/ss/project/{project_id}/asset/{asset_id}", json=payload)
+
+    def get_project_workspace(self, project_id: str) -> Any:
+        return self._get(f"/api/ss/project/{project_id}/workspace")
+
+    def upload_workspace_file(self, project_id: str, file_path: str) -> Any:
+        if not os.path.isfile(file_path):
+            raise FileNotFoundError(file_path)
+        with open(file_path, "rb") as handle:
+            return self._post_files(
+                f"/api/ss/project/{project_id}/workspace/file",
+                files={"file": (os.path.basename(file_path), handle)},
+            )
+
+    def download_workspace_file(self, project_id: str, file_name: str) -> Any:
+        return self._get(f"/api/ss/project/{project_id}/workspace/{file_name}")
+
+    def get_project_notes(self, project_id: str) -> Any:
+        return self._get(f"/api/ss/project/{project_id}/notes")
+
+    def create_project_note(self, project_id: str, payload: Dict[str, Any]) -> Any:
+        return self._post(f"/api/ss/project/{project_id}/note", json=payload)
+
+    def update_project_note(self, project_id: str, note_id: str, payload: Dict[str, Any]) -> Any:
+        return self._put(f"/api/ss/project/{project_id}/note/{note_id}", json=payload)
+
+    def create_project_workspace_note(self, project_id: str, payload: Dict[str, Any]) -> Any:
+        return self._post(f"/api/ss/project/{project_id}/workspace/note", json=payload)
+
+    def update_project_workspace_note(self, project_id: str, note_id: str, payload: Dict[str, Any]) -> Any:
+        return self._put(f"/api/ss/project/{project_id}/workspace/note/{note_id}", json=payload)
+
+    def get_project_membership_administrators(self, project_id: str) -> Any:
+        return self._get(f"/api/ss/project/{project_id}/member-admins")
+
+    def _normalize_member_admin_payload(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+        if "member_admins" in payload:
+            return payload
+        return {"member_admins": [payload]}
+
+    def add_project_membership_administrators(self, project_id: str, payload: Dict[str, Any]) -> Any:
+        return self._post(
+            f"/api/ss/project/{project_id}/member-admins",
+            json=self._normalize_member_admin_payload(payload),
+        )
+
+    def update_project_membership_administrators(self, project_id: str, payload: Dict[str, Any]) -> Any:
+        return self._put(
+            f"/api/ss/project/{project_id}/member-admins",
+            json=self._normalize_member_admin_payload(payload),
+        )
+
+    def remove_project_membership_administrators(self, project_id: str, payload: Dict[str, Any]) -> Any:
+        return self._delete(
+            f"/api/ss/project/{project_id}/member-admins",
+            json=self._normalize_member_admin_payload(payload),
+        )
+
+    def invite_user_to_project(self, project_id: str, payload: Dict[str, Any]) -> Any:
+        return self._post(f"/api/ss/project/{project_id}/invite", json=payload)
+
+    def invite_users_to_project_team(self, project_id: str, payload: Dict[str, Any]) -> Any:
+        return self._post(f"/api/ss/project/{project_id}/team/invite", json=payload)
+
+    def remove_project_team_members(self, project_id: str, payload: Dict[str, Any]) -> Any:
+        return self._put(f"/api/ss/project/{project_id}/team/remove", json=payload)
+
+    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 extract_projects_list(self, projects_data: Any) -> List[Dict[str, Any]]:
+        if not isinstance(projects_data, dict):
+            return []
+        projects = projects_data.get("projects") or (projects_data.get("data") or {}).get("projects") or []
+        return [proj for proj in projects if isinstance(proj, dict)] if isinstance(projects, list) else []
+
+    def find_project_by_name(
+        self, name: str, *, projects_data: Optional[Any] = None, case_insensitive: bool = True
+    ) -> Optional[Dict[str, Any]]:
+        data = projects_data if projects_data is not None else self.get_projects()
+        desired = self._normalize_string(name) if case_insensitive else name
+        for project in self.extract_projects_list(data):
+            project_name = project.get("project_name") or project.get("name")
+            if not isinstance(project_name, str):
+                continue
+            candidate = self._normalize_string(project_name) if case_insensitive else project_name
+            if candidate == desired:
+                return project
+        return None
+
+    def _normalize_string(self, value: str) -> str:
+        return value.strip().lower()

pyattackforge/resources/reports.py

@@ -0,0 +1,20 @@
+"""Resource: reports."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from .base import BaseResource
+
+
+class ReportsResource(BaseResource):
+    """Reports API resource wrapper."""
+
+    def get_project_report(self, project_id: str, report_type: str, params: Optional[Dict[str, Any]] = None) -> Any:
+        return self._get(f"/api/ss/project/{project_id}/report/{report_type}", params=params)
+
+    def get_project_report_data(self, project_id: str, report_type: str, payload: Dict[str, Any]) -> Any:
+        return self._post(f"/api/ss/project/{project_id}/report/{report_type}", json=payload)
+
+    def update_exec_summary_notes(self, project_id: str, payload: Dict[str, Any]) -> Any:
+        return self._put(f"/api/ss/project/{project_id}/execSummaryNotes", json=payload)

pyattackforge/resources/users.py

@@ -0,0 +1,59 @@
+"""Resource: users."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+import urllib.parse
+
+from .base import BaseResource
+
+
+class UsersResource(BaseResource):
+    """Users API resource wrapper."""
+
+    def create_user(self, payload: Dict[str, Any]) -> Any:
+        return self._post("/api/ss/user", json=payload)
+
+    def create_users(self, users: Any) -> Any:
+        return self._post("/api/ss/users", json=users)
+
+    def get_user(self, user_id: str) -> Any:
+        return self._get(f"/api/ss/users/{user_id}")
+
+    def get_users(self, params: Optional[Dict[str, Any]] = None) -> Any:
+        return self._get("/api/ss/users", params=params)
+
+    def get_user_by_email(self, email: str) -> Any:
+        email_value = urllib.parse.quote(email, safe="")
+        return self._get(f"/api/ss/users/email/{email_value}")
+
+    def get_user_by_username(self, username: str) -> Any:
+        username_value = urllib.parse.quote(username, safe="")
+        return self._get(f"/api/ss/users/username/{username_value}")
+
+    def update_user(self, user_id: str, payload: Dict[str, Any]) -> Any:
+        return self._put(f"/api/ss/user/{user_id}", json=payload)
+
+    def activate_user(self, user_id: str) -> Any:
+        return self._put(f"/api/ss/user/{user_id}/activate")
+
+    def deactivate_user(self, user_id: str) -> Any:
+        return self._put(f"/api/ss/user/{user_id}/deactivate")
+
+    def add_user_to_group(self, payload: Dict[str, Any]) -> Any:
+        return self._post("/api/ss/group/user", json=payload)
+
+    def update_user_access_on_group(self, user_id: str, payload: Dict[str, Any]) -> Any:
+        return self._put(f"/api/ss/group/user/{user_id}", json=payload)
+
+    def get_user_groups(self, user_id: str) -> Any:
+        return self._get(f"/api/ss/user/{user_id}/groups")
+
+    def get_user_projects(self, user_id: str) -> Any:
+        return self._get(f"/api/ss/user/{user_id}/projects")
+
+    def get_user_audit_logs(self, user_id: str, params: Optional[Dict[str, Any]] = None) -> Any:
+        return self._get(f"/api/ss/user/{user_id}/auditlogs", params=params)
+
+    def get_user_login_history(self, user_id: str, params: Optional[Dict[str, Any]] = None) -> Any:
+        return self._get(f"/api/ss/user/{user_id}/logins", params=params)

pyattackforge/resources/writeups.py

@@ -0,0 +1,79 @@
+"""Resource: writeups."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+import os
+
+from ..cache import TTLCache
+from .base import BaseResource
+
+
+class WriteupsResource(BaseResource):
+    """Writeups (library vulnerabilities) API resource wrapper."""
+
+    def __init__(self, transport) -> None:  # type: ignore[override]
+        super().__init__(transport)
+        self._cache = TTLCache(default_ttl=300.0)
+
+    def get_writeups(self, params: Optional[Dict[str, Any]] = None, *, force_refresh: bool = False) -> Any:
+        cache_key = "writeups:list"
+        if params is None and not force_refresh:
+            cached = self._cache.get(cache_key)
+            if cached is not None:
+                return cached
+        data = self._get("/api/ss/library", params=params)
+        if params is None:
+            self._cache.set(cache_key, data)
+        return data
+
+    def get_writeup_files(self, *, writeup_id: Optional[str] = None, name: Optional[str] = None) -> list[Dict[str, Any]]:
+        """
+        Fetch file entries for a writeup (library vulnerability).
+
+        The SSAPI returns file storage names under `files[].storage_name` when
+        querying `/api/ss/library` with `id` or `name` filters.
+        """
+        params: Dict[str, Any] = {}
+        if writeup_id:
+            params["id"] = writeup_id
+        if name:
+            params["name"] = name
+        data = self.get_writeups(params=params, force_refresh=True)
+        if isinstance(data, dict):
+            candidates = data.get("vulnerabilities") or data.get("library") or data.get("issues") or []
+        elif isinstance(data, list):
+            candidates = data
+        else:
+            candidates = []
+        if not isinstance(candidates, list):
+            return []
+        for entry in candidates:
+            if not isinstance(entry, dict):
+                continue
+            if writeup_id and entry.get("id") != writeup_id:
+                continue
+            if name and entry.get("title") != name:
+                continue
+            files = entry.get("files")
+            if isinstance(files, list):
+                return files
+        return []
+
+    def create_writeup(self, payload: Dict[str, Any]) -> Any:
+        return self._post("/api/ss/library/vulnerability", json=payload)
+
+    def update_writeup(self, writeup_id: str, payload: Dict[str, Any]) -> Any:
+        return self._put(f"/api/ss/library/{writeup_id}", json=payload)
+
+    def upload_writeup_file(self, writeup_id: str, file_path: str) -> Any:
+        if not os.path.isfile(file_path):
+            raise FileNotFoundError(file_path)
+        with open(file_path, "rb") as handle:
+            return self._post_files(
+                f"/api/ss/library/{writeup_id}/file",
+                files={"file": (os.path.basename(file_path), handle)},
+            )
+
+    def download_writeup_file(self, writeup_id: str, file_name: str) -> Any:
+        return self._get(f"/api/ss/library/{writeup_id}/file/{file_name}")

pyattackforge/transport.py

@@ -0,0 +1,134 @@
+"""HTTP transport for AttackForge SSAPI."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+import time
+import httpx
+
+from .exceptions import APIError
+from .config import ClientConfig
+
+
+@dataclass
+class TransportResponse:
+    status_code: int
+    data: Any
+    headers: Dict[str, str]
+
+
+class AttackForgeTransport:
+    def __init__(self, config: ClientConfig, client: Optional[httpx.Client] = None) -> None:
+        self._config = config
+        self._client = client or httpx.Client(
+            base_url=config.base_url.rstrip("/"),
+            timeout=config.timeout,
+            headers={
+                "X-SSAPI-KEY": config.api_key,
+                "Connection": "close",
+                "User-Agent": config.user_agent,
+            },
+            http2=config.http2,
+        )
+        self._http2_client: Optional[httpx.Client] = None
+
+    def close(self) -> None:
+        self._client.close()
+        if self._http2_client is not None:
+            self._http2_client.close()
+
+    def _get_http2_client(self) -> httpx.Client:
+        if self._http2_client is None:
+            self._http2_client = httpx.Client(
+                base_url=self._config.base_url.rstrip("/"),
+                timeout=self._config.timeout,
+                headers={
+                    "X-SSAPI-KEY": self._config.api_key,
+                    "Connection": "close",
+                    "User-Agent": self._config.user_agent,
+                },
+                http2=True,
+            )
+        return self._http2_client
+
+    def request(
+        self,
+        method: str,
+        path: str,
+        *,
+        params: Optional[Dict[str, Any]] = None,
+        json: Optional[Dict[str, Any]] = None,
+        files: Optional[Dict[str, Any]] = None,
+        data: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+        timeout: Optional[float] = None,
+        retries: Optional[int] = None,
+    ) -> TransportResponse:
+        attempt = 0
+        max_retries = self._config.max_retries if retries is None else retries
+        backoff = self._config.backoff_factor
+        while True:
+            try:
+                request_headers = None
+                if headers:
+                    request_headers = headers
+                if files:
+                    request_headers = dict(self._client.headers)
+                    request_headers.pop("Content-Type", None)
+                    if headers:
+                        request_headers.update(headers)
+                response = self._client.request(
+                    method=method,
+                    url=path,
+                    params=params,
+                    json=json,
+                    files=files,
+                    data=data,
+                    headers=request_headers,
+                    timeout=timeout,
+                )
+            except httpx.HTTPError as exc:
+                if attempt >= max_retries:
+                    raise APIError(status_code=0, message=str(exc)) from exc
+                time.sleep(backoff * (2 ** attempt))
+                attempt += 1
+                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 response.status_code in (429, 500, 502, 503, 504) and attempt < max_retries:
+                time.sleep(backoff * (2 ** attempt))
+                attempt += 1
+                continue
+
+            data_out: Any
+            if response.headers.get("content-type", "").startswith("application/json"):
+                try:
+                    data_out = response.json()
+                except ValueError:
+                    data_out = response.text
+            else:
+                data_out = response.content
+
+            if response.status_code >= 400:
+                raise APIError(response.status_code, response.text, payload={"data": data_out})
+
+            return TransportResponse(
+                status_code=response.status_code,
+                data=data_out,
+                headers=dict(response.headers),
+            )