"""Evidence transfer for the assessment import/export bundle flows.
These are module-level functions rather than :class:`~trustwise.sdk.platform.client.AssessmentResource` methods on
purpose. The CLI auto-generates a Click command for every public method it finds on
a resource class, and a method taking raw ``bytes`` would turn into a nonsense
subcommand. Keeping them here makes them invisible to that machinery.
The step order they serve is the same one the portal implements, so both clients
behave identically against the same server:
1. ``POST /assessments/import`` returns presigned uploads plus a ready-to-PUT draft.
2. Upload every evidence file in parallel (:func:`upload_evidence`).
3. ``PUT /assessments/{id}/answers`` with the echoed draft (:func:`answers_payload`).
4. Re-upload anything the server rejected, then ``POST /assessments/{id}/submit``.
Export is the mirror: ``GET /assessments/{id}/export`` then
:func:`download_evidence`.
"""
from __future__ import annotations
from collections.abc import Iterable
from datetime import datetime
from typing import Any, NamedTuple
from pydantic import BaseModel
from trustwise.sdk.exceptions import TrustwiseAPIError, TrustwiseValidationError
from trustwise.sdk.utils.blob import (
DEFAULT_CONCURRENCY,
get_presigned,
map_parallel,
put_presigned,
same_origin,
)
[docs]
class EvidenceBlob(NamedTuple):
"""One evidence file's bytes plus the metadata an upload needs.
Attributes:
data: Raw file bytes.
filename: Original filename, sent as ``x-filename`` on a loopback upload
(the local backend falls back to the file id without it).
content_type: Content type, when the bundle recorded one.
"""
data: bytes
filename: str
content_type: str | None = None
[docs]
def _as_dict(value: Any) -> dict[str, Any]:
"""Normalize a pydantic model or mapping into a plain dict.
Responses arrive as generated models when the spec is current and as plain
dicts otherwise (and in tests), so every accessor below works on one shape.
Args:
value: A pydantic model or a mapping.
Returns:
A plain dict view of the value.
"""
if isinstance(value, BaseModel):
return value.model_dump(mode="json")
return dict(value)
[docs]
def answers_payload(draft_answers: Any) -> dict[str, Any]:
"""Render the import response's ``draft_answers`` for ``PUT /answers``.
The server hands back a ready-to-PUT document and expects it echoed unchanged.
This is that echo — JSON-rendered because ``requests`` cannot serialize pydantic
models, with unset fields dropped so nulls the server never sent (``object_key``)
stay off the wire.
Args:
draft_answers: ``draft_answers`` from the import response.
Returns:
The same document as plain JSON-ready dicts, keyed by question path.
"""
if isinstance(draft_answers, BaseModel):
return draft_answers.model_dump(mode="json", exclude_none=True)
return {
path: entry.model_dump(mode="json", exclude_none=True)
if isinstance(entry, BaseModel)
else entry
for path, entry in dict(draft_answers).items()
}
[docs]
def answers_body(draft_answers: Any, updated_at: Any) -> dict[str, Any]:
"""Build the ``PUT /answers`` body from an import response.
Args:
draft_answers: ``draft_answers`` from the import response, echoed as-is.
updated_at: The optimistic-lock token from the same response. Accepts a
``datetime`` (which the generated models produce and ``requests`` cannot
serialize) or an already-formatted string.
Returns:
The request body for ``PUT /assessments/{id}/answers``.
Raises:
TrustwiseValidationError: If no lock token is available. The server rejects a
null one, so failing here beats sending an empty string.
"""
if updated_at is None or updated_at == "":
raise TrustwiseValidationError(
"Import response carried no 'updated_at' to lock the answers write against"
)
token = updated_at.isoformat() if isinstance(updated_at, datetime) else str(updated_at)
return {
"expected_updated_at": token,
"draft_answers": answers_payload(draft_answers),
}
[docs]
def upload_evidence(
uploads: Iterable[Any],
payloads: dict[str, EvidenceBlob],
*,
api_base_url: str,
api_key: str,
max_workers: int = DEFAULT_CONCURRENCY,
) -> None:
"""Upload every evidence file to the URL the import response minted for it.
Uploads run in parallel with no per-file retry: the first failure aborts the
batch, and recovery is re-running the import (which mints fresh URLs for
everything). This matches the portal.
Two kinds of URL come back, and they need opposite treatment:
* A real presigned URL on the blob store's own host gets **only** the headers
the server signed. Sending a Trustwise credential there would leak it.
* A URL on the API's own origin is the local backend's loopback route, which
cannot presign and authenticates normally. That one needs the bearer token,
plus ``x-filename`` so the stored file keeps its name.
Args:
uploads: ``uploads`` from the import response.
payloads: Evidence bytes keyed by the bundle's own file id, which the
server echoes back as ``source_file_id``.
api_base_url: The Trustwise API base URL, used to recognise loopback URLs.
api_key: Trustwise API key, sent only to a loopback URL.
max_workers: Upper bound on concurrent uploads.
Raises:
TrustwiseValidationError: If an upload references bytes the bundle lacks.
TrustwiseAPIError: If any upload fails.
"""
entries = [_as_dict(upload) for upload in uploads]
for entry in entries:
if entry.get("source_file_id") not in payloads:
raise TrustwiseValidationError(
f"Bundle is missing the evidence payload for {entry.get('source_file_id')!r}"
)
def send(entry: dict[str, Any]) -> None:
blob = payloads[entry["source_file_id"]]
headers = dict(entry.get("upload_headers") or {})
url = entry["upload_url"]
if same_origin(url, api_base_url):
headers.setdefault("Content-Type", blob.content_type or "application/octet-stream")
headers["Authorization"] = f"Bearer {api_key}"
headers["x-filename"] = blob.filename
put_presigned(url, blob.data, headers)
failures = [
(entry, exc)
for entry, _, exc in map_parallel(send, entries, max_workers=max_workers)
if exc is not None
]
if failures:
detail = ", ".join(
f"{entry.get('file_id')} (HTTP {getattr(exc, 'status_code', '?')})"
for entry, exc in failures
)
raise TrustwiseAPIError(
message=f"Evidence upload failed for {len(failures)} file(s): {detail}",
response=detail,
status_code=getattr(failures[0][1], "status_code", 0),
)
[docs]
def download_evidence(
evidence: Iterable[Any],
*,
max_workers: int = DEFAULT_CONCURRENCY,
) -> dict[str, bytes]:
"""Download every evidence file listed in an export manifest.
Downloads run in parallel and send no credentials — the presigned URL is the
credential. Checksums are not verified here; the caller verifies them against
the bundle's own ``answers.json`` before writing anything.
Args:
evidence: ``evidence`` from the export response.
max_workers: Upper bound on concurrent downloads.
Returns:
Raw bytes keyed by ``file_id``.
Raises:
TrustwiseAPIError: If any download fails.
"""
entries = [_as_dict(item) for item in evidence]
results = map_parallel(
lambda entry: get_presigned(entry["download_url"]),
entries,
max_workers=max_workers,
)
failures = [(entry, exc) for entry, _, exc in results if exc is not None]
if failures:
detail = ", ".join(
f"{entry.get('filename') or entry.get('file_id')} "
f"(HTTP {getattr(exc, 'status_code', '?')})"
for entry, exc in failures
)
raise TrustwiseAPIError(
message=f"Evidence download failed for {len(failures)} file(s): {detail}",
response=detail,
status_code=getattr(failures[0][1], "status_code", 0),
)
return {entry["file_id"]: data for entry, data, _ in results if data is not None}