| 1 | import base64 |
| 2 | from io import BytesIO |
| 3 | import os |
| 4 | from pathlib import Path |
| 5 | import tempfile |
| 6 | import zipfile |
| 7 | |
| 8 | from flask import Response |
| 9 | |
| 10 | from helpers.api import ApiHandler, Input, Output, Request |
| 11 | from helpers import runtime |
| 12 | from helpers.localization import Localization |
| 13 | from api.download_work_dir_file import fetch_file, stream_file_download |
| 14 | |
| 15 | |
| 16 | class DownloadFiles(ApiHandler): |
| 17 | async def process(self, input: Input, request: Request) -> Output: |
| 18 | try: |
| 19 | paths = normalize_paths(input.get("paths", [])) |
| 20 | except ValueError as exc: |
| 21 | return Response(str(exc), status=400) |
| 22 | |
| 23 | current_path = input.get("currentPath", "") |
| 24 | |
| 25 | if not paths: |
| 26 | return Response("No file paths provided", status=400) |
| 27 | |
| 28 | try: |
| 29 | zip_file = await runtime.call_development_function( |
| 30 | create_selected_zip, paths, current_path |
| 31 | ) |
| 32 | except ValueError as exc: |
| 33 | return Response(str(exc), status=400) |
| 34 | except FileNotFoundError as exc: |
| 35 | return Response(str(exc), status=404) |
| 36 | |
| 37 | download_name = selected_archive_name(len(paths)) |
| 38 | if runtime.is_development(): |
| 39 | b64 = await runtime.call_development_function(fetch_file, zip_file) |
| 40 | file_data = BytesIO(base64.b64decode(b64)) |
| 41 | return stream_file_download(file_data, download_name=download_name) |
| 42 | |
| 43 | return stream_file_download(zip_file, download_name=download_name) |
| 44 | |
| 45 | |
| 46 | def normalize_paths(paths) -> list[str]: |
| 47 | if not isinstance(paths, list): |
| 48 | raise ValueError("Paths must be a list") |
| 49 | |
| 50 | normalized: list[str] = [] |
| 51 | seen: set[str] = set() |
| 52 | for raw_path in paths: |
| 53 | if not isinstance(raw_path, str): |
| 54 | continue |
| 55 | path = raw_path.strip() |
| 56 | if not path: |
| 57 | continue |
| 58 | if not path.startswith("/"): |
| 59 | path = f"/{path}" |
| 60 | if path not in seen: |
| 61 | normalized.append(path) |
| 62 | seen.add(path) |
| 63 | |
| 64 | return normalized |
| 65 | |
| 66 | |
| 67 | def selected_archive_name(count: int) -> str: |
| 68 | stamp = Localization.get().now().strftime("%Y%m%d-%H%M%S") |
| 69 | return f"agent-zero-selected-{count}-{stamp}.zip" |
| 70 | |
| 71 | |
| 72 | def create_selected_zip(paths: list[str], current_path: str = "") -> str: |
| 73 | base_dir = Path("/") |
| 74 | current_dir = resolve_download_path(current_path, base_dir) if current_path else None |
| 75 | if current_dir and current_dir.is_file(): |
| 76 | current_dir = current_dir.parent |
| 77 | |
| 78 | selected_paths = [] |
| 79 | for path in normalize_paths(paths): |
| 80 | resolved = resolve_download_path(path, base_dir) |
| 81 | if resolved.exists(): |
| 82 | selected_paths.append(resolved) |
| 83 | |
| 84 | selected_paths = collapse_nested_paths(selected_paths) |
| 85 | if not selected_paths: |
| 86 | raise FileNotFoundError("No selected files were found") |
| 87 | |
| 88 | zip_file_path = tempfile.NamedTemporaryFile(suffix=".zip", delete=False).name |
| 89 | used_names: set[str] = set() |
| 90 | |
| 91 | with zipfile.ZipFile( |
| 92 | zip_file_path, "w", compression=zipfile.ZIP_DEFLATED, allowZip64=True |
| 93 | ) as zip_file: |
| 94 | for source_path in selected_paths: |
| 95 | arc_root = unique_archive_name( |
| 96 | archive_root_name(source_path, current_dir, base_dir), used_names |
| 97 | ) |
| 98 | write_zip_entry(zip_file, source_path, arc_root) |
| 99 | |
| 100 | return zip_file_path |
| 101 | |
| 102 | |
| 103 | def resolve_download_path(path: str, base_dir: Path) -> Path: |
| 104 | if not path: |
| 105 | raise ValueError("Invalid file path") |
| 106 | |
| 107 | candidate = Path(path) |
| 108 | resolved = candidate.resolve() if candidate.is_absolute() else (base_dir / candidate).resolve() |
| 109 | |
| 110 | try: |
| 111 | resolved.relative_to(base_dir) |
| 112 | except ValueError as exc: |
| 113 | raise ValueError("Invalid file path") from exc |
| 114 | |
| 115 | return resolved |
| 116 | |
| 117 | |
| 118 | def collapse_nested_paths(paths: list[Path]) -> list[Path]: |
| 119 | collapsed: list[Path] = [] |
| 120 | for path in sorted(paths, key=lambda item: len(item.parts)): |
| 121 | if any(path == parent or parent in path.parents for parent in collapsed): |
| 122 | continue |
| 123 | collapsed.append(path) |
| 124 | return collapsed |
| 125 | |
| 126 | |
| 127 | def archive_root_name(source_path: Path, current_dir: Path | None, base_dir: Path) -> str: |
| 128 | if current_dir: |
| 129 | try: |
| 130 | return source_path.relative_to(current_dir).as_posix().strip("/") |
| 131 | except ValueError: |
| 132 | pass |
| 133 | |
| 134 | try: |
| 135 | return source_path.relative_to(base_dir).as_posix().strip("/") |
| 136 | except ValueError: |
| 137 | return source_path.name |
| 138 | |
| 139 | |
| 140 | def unique_archive_name(name: str, used_names: set[str]) -> str: |
| 141 | clean_name = name or "selection" |
| 142 | if clean_name not in used_names: |
| 143 | used_names.add(clean_name) |
| 144 | return clean_name |
| 145 | |
| 146 | stem, suffix = os.path.splitext(clean_name) |
| 147 | index = 2 |
| 148 | while True: |
| 149 | candidate = f"{stem}-{index}{suffix}" |
| 150 | if candidate not in used_names: |
| 151 | used_names.add(candidate) |
| 152 | return candidate |
| 153 | index += 1 |
| 154 | |
| 155 | |
| 156 | def write_zip_entry(zip_file: zipfile.ZipFile, source_path: Path, arc_root: str) -> None: |
| 157 | if source_path.is_dir(): |
| 158 | wrote_any = False |
| 159 | for root, dirs, file_names in os.walk(source_path): |
| 160 | dirs.sort() |
| 161 | file_names.sort() |
| 162 | root_path = Path(root) |
| 163 | rel_root = root_path.relative_to(source_path) |
| 164 | |
| 165 | if not dirs and not file_names: |
| 166 | empty_dir = Path(arc_root) / rel_root |
| 167 | zip_file.writestr(empty_dir.as_posix().rstrip("/") + "/", "") |
| 168 | |
| 169 | for file_name in file_names: |
| 170 | file_path = root_path / file_name |
| 171 | rel_path = file_path.relative_to(source_path) |
| 172 | zip_file.write(file_path, (Path(arc_root) / rel_path).as_posix()) |
| 173 | wrote_any = True |
| 174 | |
| 175 | if not wrote_any: |
| 176 | zip_file.writestr(Path(arc_root).as_posix().rstrip("/") + "/", "") |
| 177 | return |
| 178 | |
| 179 | zip_file.write(source_path, arc_root) |