| 1 | import base64 |
| 2 | import os |
| 3 | from pathlib import Path |
| 4 | from urllib.parse import quote |
| 5 | from helpers.api import ApiHandler, Request, Response, send_file |
| 6 | from helpers import files, runtime |
| 7 | import io |
| 8 | from mimetypes import guess_type |
| 9 | |
| 10 | |
| 11 | IMAGE_EXTENSIONS = ( |
| 12 | ".jpg", |
| 13 | ".jpeg", |
| 14 | ".png", |
| 15 | ".gif", |
| 16 | ".bmp", |
| 17 | ".webp", |
| 18 | ".svg", |
| 19 | ".ico", |
| 20 | ".svgz", |
| 21 | ) |
| 22 | SVG_EXTENSIONS = (".svg", ".svgz") |
| 23 | SVG_CONTENT_SECURITY_POLICY = ( |
| 24 | "sandbox; default-src 'none'; script-src 'none'; " |
| 25 | "img-src 'self' data:; style-src 'unsafe-inline'" |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | class ImageGet(ApiHandler): |
| 30 | |
| 31 | @classmethod |
| 32 | def get_methods(cls) -> list[str]: |
| 33 | return ["GET"] |
| 34 | |
| 35 | async def process(self, input: dict, request: Request) -> dict | Response: |
| 36 | # input data |
| 37 | path = input.get("path", request.args.get("path", "")) |
| 38 | |
| 39 | if not path: |
| 40 | raise ValueError("No path provided") |
| 41 | |
| 42 | # get file extension and info |
| 43 | file_ext = os.path.splitext(path)[1].lower() |
| 44 | filename = os.path.basename(path) |
| 45 | |
| 46 | if file_ext in IMAGE_EXTENSIONS: |
| 47 | try: |
| 48 | local_path = _resolve_allowed_image_path(path) |
| 49 | except ValueError as exc: |
| 50 | return Response(str(exc), status=403, mimetype="text/plain") |
| 51 | |
| 52 | # in development environment, try to serve the image from local file system if exists, otherwise from docker |
| 53 | if runtime.is_development(): |
| 54 | if files.exists(local_path): |
| 55 | response = send_file(local_path) |
| 56 | else: |
| 57 | # Try fetching from Docker via RFC as fallback |
| 58 | try: |
| 59 | remote_path = await runtime.call_development_function( |
| 60 | _resolve_allowed_image_path, path |
| 61 | ) |
| 62 | if await runtime.call_development_function( |
| 63 | files.exists, remote_path |
| 64 | ): |
| 65 | b64_content = await runtime.call_development_function( |
| 66 | files.read_file_base64, remote_path |
| 67 | ) |
| 68 | file_content = base64.b64decode(b64_content) |
| 69 | mime_type, _ = guess_type(filename) |
| 70 | if not mime_type: |
| 71 | mime_type = "application/octet-stream" |
| 72 | response = send_file( |
| 73 | io.BytesIO(file_content), |
| 74 | mimetype=mime_type, |
| 75 | as_attachment=False, |
| 76 | download_name=filename, |
| 77 | ) |
| 78 | else: |
| 79 | response = _send_fallback_icon("image") |
| 80 | except Exception: |
| 81 | response = _send_fallback_icon("image") |
| 82 | else: |
| 83 | if files.exists(local_path): |
| 84 | response = send_file(local_path) |
| 85 | else: |
| 86 | response = _send_fallback_icon("image") |
| 87 | |
| 88 | _set_image_headers(response, filename, file_ext) |
| 89 | return response |
| 90 | else: |
| 91 | # Handle non-image files with fallback icons |
| 92 | return _send_file_type_icon(file_ext, filename) |
| 93 | |
| 94 | |
| 95 | def _resolve_allowed_image_path(path: str) -> str: |
| 96 | """Resolve a requested image path and keep it inside Agent Zero's base dir.""" |
| 97 | |
| 98 | if runtime.is_development(): |
| 99 | candidate = Path(files.fix_dev_path(path)) |
| 100 | else: |
| 101 | candidate = Path(files.get_abs_path(path)) |
| 102 | |
| 103 | if not candidate.is_absolute(): |
| 104 | candidate = Path(files.get_base_dir()) / candidate |
| 105 | |
| 106 | base_dir = Path(files.get_base_dir()).resolve() |
| 107 | resolved = candidate.resolve(strict=False) |
| 108 | |
| 109 | try: |
| 110 | resolved.relative_to(base_dir) |
| 111 | except ValueError as exc: |
| 112 | raise ValueError("Path is outside of allowed directory") from exc |
| 113 | |
| 114 | return str(resolved) |
| 115 | |
| 116 | |
| 117 | def _set_image_headers(response: Response, filename: str, file_ext: str) -> None: |
| 118 | # Add cache headers for better device sync performance. |
| 119 | response.headers["Cache-Control"] = "public, max-age=3600" |
| 120 | response.headers["X-File-Type"] = "image" |
| 121 | response.headers["X-File-Name"] = quote(filename) |
| 122 | response.headers["X-Content-Type-Options"] = "nosniff" |
| 123 | if file_ext in SVG_EXTENSIONS: |
| 124 | response.headers["Content-Security-Policy"] = SVG_CONTENT_SECURITY_POLICY |
| 125 | |
| 126 | |
| 127 | def _send_file_type_icon(file_ext, filename=None): |
| 128 | """Return appropriate icon for file type""" |
| 129 | |
| 130 | # Map file extensions to icon names |
| 131 | icon_mapping = { |
| 132 | # Archive files |
| 133 | ".zip": "archive", |
| 134 | ".rar": "archive", |
| 135 | ".7z": "archive", |
| 136 | ".tar": "archive", |
| 137 | ".gz": "archive", |
| 138 | # Document files |
| 139 | ".pdf": "document", |
| 140 | ".doc": "document", |
| 141 | ".docx": "document", |
| 142 | ".txt": "document", |
| 143 | ".rtf": "document", |
| 144 | ".odt": "document", |
| 145 | # Code files |
| 146 | ".py": "code", |
| 147 | ".js": "code", |
| 148 | ".html": "code", |
| 149 | ".css": "code", |
| 150 | ".json": "code", |
| 151 | ".xml": "code", |
| 152 | ".md": "code", |
| 153 | ".yml": "code", |
| 154 | ".yaml": "code", |
| 155 | ".sql": "code", |
| 156 | ".sh": "code", |
| 157 | ".bat": "code", |
| 158 | # Spreadsheet files |
| 159 | ".xls": "document", |
| 160 | ".xlsx": "document", |
| 161 | ".csv": "document", |
| 162 | # Presentation files |
| 163 | ".ppt": "document", |
| 164 | ".pptx": "document", |
| 165 | ".odp": "document", |
| 166 | } |
| 167 | |
| 168 | # Get icon name, default to 'file' if not found |
| 169 | icon_name = icon_mapping.get(file_ext, "file") |
| 170 | |
| 171 | response = _send_fallback_icon(icon_name) |
| 172 | |
| 173 | # Add headers for device sync |
| 174 | if hasattr(response, "headers"): |
| 175 | response.headers["Cache-Control"] = ( |
| 176 | "public, max-age=86400" # Cache icons for 24 hours |
| 177 | ) |
| 178 | response.headers["X-File-Type"] = "icon" |
| 179 | response.headers["X-Icon-Type"] = icon_name |
| 180 | if filename: |
| 181 | response.headers["X-File-Name"] = quote(filename) |
| 182 | |
| 183 | return response |
| 184 | |
| 185 | |
| 186 | def _send_fallback_icon(icon_name): |
| 187 | """Return fallback icon from public directory""" |
| 188 | |
| 189 | # Path to public icons |
| 190 | icon_path = files.get_abs_path(f"webui/public/{icon_name}.svg") |
| 191 | |
| 192 | # Check if specific icon exists, fallback to generic file icon |
| 193 | if not os.path.exists(icon_path): |
| 194 | icon_path = files.get_abs_path("webui/public/file.svg") |
| 195 | |
| 196 | # Final fallback if file.svg doesn't exist |
| 197 | if not os.path.exists(icon_path): |
| 198 | raise ValueError(f"Fallback icon not found: {icon_path}") |
| 199 | |
| 200 | return send_file(icon_path, mimetype="image/svg+xml") |