| 1 | from __future__ import annotations |
| 2 | |
| 3 | import base64 |
| 4 | import binascii |
| 5 | import mimetypes |
| 6 | import uuid |
| 7 | from dataclasses import dataclass |
| 8 | from pathlib import Path |
| 9 | from urllib.parse import urlparse |
| 10 | |
| 11 | from helpers import files |
| 12 | |
| 13 | |
| 14 | DEFAULT_MAX_ARTIFACT_SIZE_BYTES = 25 * 1024 * 1024 |
| 15 | |
| 16 | |
| 17 | class MediaArtifactError(ValueError): |
| 18 | pass |
| 19 | |
| 20 | |
| 21 | class EmptyBase64Data(MediaArtifactError): |
| 22 | pass |
| 23 | |
| 24 | |
| 25 | class InvalidBase64Data(MediaArtifactError): |
| 26 | pass |
| 27 | |
| 28 | |
| 29 | class ArtifactTooLarge(MediaArtifactError): |
| 30 | def __init__(self, size: int, limit: int): |
| 31 | super().__init__(f"artifact is too large ({size} bytes, limit {limit} bytes)") |
| 32 | self.size = size |
| 33 | self.limit = limit |
| 34 | |
| 35 | |
| 36 | @dataclass(frozen=True) |
| 37 | class Base64Payload: |
| 38 | data: str |
| 39 | payload: bytes |
| 40 | size: int |
| 41 | |
| 42 | |
| 43 | @dataclass(frozen=True) |
| 44 | class ImageDataUrl: |
| 45 | url: str |
| 46 | mime: str |
| 47 | size: int |
| 48 | |
| 49 | |
| 50 | @dataclass(frozen=True) |
| 51 | class SavedArtifact: |
| 52 | path: str |
| 53 | mime: str |
| 54 | size: int |
| 55 | |
| 56 | |
| 57 | def compact_base64(data: str) -> str: |
| 58 | return "".join(char for char in str(data or "") if not char.isspace()) |
| 59 | |
| 60 | |
| 61 | def estimated_base64_decoded_size(data: str) -> int: |
| 62 | compact_length = len(compact_base64(data)) |
| 63 | return (compact_length * 3) // 4 |
| 64 | |
| 65 | |
| 66 | def decode_base64_payload( |
| 67 | data: str, |
| 68 | *, |
| 69 | max_bytes: int | None = None, |
| 70 | ) -> Base64Payload: |
| 71 | compact = compact_base64(data) |
| 72 | if not compact: |
| 73 | raise EmptyBase64Data("base64 data is empty") |
| 74 | |
| 75 | if max_bytes is not None: |
| 76 | estimated_size = estimated_base64_decoded_size(compact) |
| 77 | if estimated_size > max_bytes: |
| 78 | raise ArtifactTooLarge(estimated_size, max_bytes) |
| 79 | |
| 80 | try: |
| 81 | payload = base64.b64decode(compact, validate=True) |
| 82 | except (binascii.Error, ValueError) as exc: |
| 83 | raise InvalidBase64Data("base64 data could not be decoded") from exc |
| 84 | |
| 85 | if max_bytes is not None and len(payload) > max_bytes: |
| 86 | raise ArtifactTooLarge(len(payload), max_bytes) |
| 87 | |
| 88 | return Base64Payload(data=compact, payload=payload, size=len(payload)) |
| 89 | |
| 90 | |
| 91 | def normalize_mime( |
| 92 | mime_type: str, |
| 93 | *, |
| 94 | default: str = "application/octet-stream", |
| 95 | required_prefix: str = "", |
| 96 | ) -> str: |
| 97 | value = str(mime_type or "").strip().lower() |
| 98 | if not value: |
| 99 | return default |
| 100 | if required_prefix and not value.startswith(required_prefix): |
| 101 | return default |
| 102 | return value |
| 103 | |
| 104 | |
| 105 | def guess_extension(mime_type: str, fallback: str = ".bin") -> str: |
| 106 | ext = mimetypes.guess_extension(str(mime_type or "").strip().lower()) or fallback |
| 107 | return ".jpg" if ext == ".jpe" else ext |
| 108 | |
| 109 | |
| 110 | def filename_from_uri(uri: str) -> str: |
| 111 | value = str(uri or "").strip() |
| 112 | if not value: |
| 113 | return "" |
| 114 | parsed = urlparse(value) |
| 115 | return Path(parsed.path or value).name |
| 116 | |
| 117 | |
| 118 | def safe_filename( |
| 119 | value: str, |
| 120 | *, |
| 121 | default: str = "artifact.bin", |
| 122 | default_extension: str = ".bin", |
| 123 | ) -> str: |
| 124 | source = str(value or "").strip() or default |
| 125 | cleaned = "".join( |
| 126 | char if char.isalnum() or char in {"-", "_", "."} else "_" |
| 127 | for char in source |
| 128 | ) |
| 129 | cleaned = cleaned.strip("._") or default |
| 130 | if "." not in cleaned: |
| 131 | ext = default_extension if default_extension.startswith(".") else f".{default_extension}" |
| 132 | cleaned += ext |
| 133 | return cleaned |
| 134 | |
| 135 | |
| 136 | def image_data_url_from_base64( |
| 137 | data: str, |
| 138 | *, |
| 139 | mime_type: str = "image/png", |
| 140 | max_bytes: int | None = None, |
| 141 | ) -> ImageDataUrl: |
| 142 | payload = decode_base64_payload(data, max_bytes=max_bytes) |
| 143 | safe_mime = normalize_mime( |
| 144 | mime_type, |
| 145 | default="image/png", |
| 146 | required_prefix="image/", |
| 147 | ) |
| 148 | return ImageDataUrl( |
| 149 | url=f"data:{safe_mime};base64,{payload.data}", |
| 150 | mime=safe_mime, |
| 151 | size=payload.size, |
| 152 | ) |
| 153 | |
| 154 | |
| 155 | def save_base64_artifact( |
| 156 | data: str, |
| 157 | *, |
| 158 | mime_type: str = "application/octet-stream", |
| 159 | directory_parts: tuple[str, ...], |
| 160 | preferred_name: str = "", |
| 161 | default_filename: str = "artifact.bin", |
| 162 | max_bytes: int | None = None, |
| 163 | ) -> SavedArtifact: |
| 164 | payload = decode_base64_payload(data, max_bytes=max_bytes) |
| 165 | safe_mime = normalize_mime(mime_type) |
| 166 | preferred_filename = filename_from_uri(preferred_name) or default_filename |
| 167 | default_extension = guess_extension(safe_mime, Path(default_filename).suffix or ".bin") |
| 168 | filename = safe_filename( |
| 169 | preferred_filename, |
| 170 | default=default_filename, |
| 171 | default_extension=default_extension, |
| 172 | ) |
| 173 | filename_path = Path(filename) |
| 174 | stem = filename_path.stem or Path(default_filename).stem or "artifact" |
| 175 | suffix = filename_path.suffix or default_extension |
| 176 | |
| 177 | artifact_dir = Path(files.get_abs_path(*directory_parts)) |
| 178 | artifact_dir.mkdir(parents=True, exist_ok=True) |
| 179 | path = artifact_dir / f"{stem}_{uuid.uuid4().hex[:8]}{suffix}" |
| 180 | path.write_bytes(payload.payload) |
| 181 | |
| 182 | return SavedArtifact( |
| 183 | path=files.normalize_a0_path(str(path)), |
| 184 | mime=safe_mime, |
| 185 | size=payload.size, |
| 186 | ) |