main
py 370 lines 13.2 KB
Raw
1 from __future__ import annotations
2
3 from datetime import datetime
4 import json
5 import os
6 import time
7 import urllib.request
8 import uuid
9 import zipfile
10 from pathlib import Path
11 from typing import Any
12
13 from helpers import files, print_style, plugins, git
14 from helpers.localization import Localization
15 from helpers import yaml as yaml_helper
16 from helpers.plugins import (
17 META_FILE_NAME,
18 PluginMetadata,
19 get_plugins_list,
20 after_plugin_change,
21 )
22 from werkzeug.datastructures import FileStorage
23 from werkzeug.utils import secure_filename
24
25
26 def _get_user_plugins_dir() -> str:
27 """Return absolute path to usr/plugins/."""
28 return files.get_abs_path(files.USER_DIR, files.PLUGINS_DIR)
29
30
31 def _get_plugin_name(meta: PluginMetadata) -> str:
32 plugin_name = (meta.name or "").strip()
33 if not plugin_name:
34 raise ValueError(f"{META_FILE_NAME} is missing required field 'name'")
35 return plugin_name
36
37
38 def validate_plugin_dir(path: str, plugin_name: str = "") -> PluginMetadata:
39 """Check directory contains plugin.yaml and return parsed metadata.
40 Raises ValueError if plugin.yaml is missing or invalid."""
41 meta_path = os.path.join(path, META_FILE_NAME)
42 if not os.path.isfile(meta_path):
43 raise ValueError(f"No {META_FILE_NAME} found in {os.path.basename(path)}")
44 with open(meta_path, "r", encoding="utf-8") as f:
45 content = f.read()
46 data = yaml_helper.loads(content)
47 model = PluginMetadata.model_validate(data)
48 if plugin_name and plugin_name != model.name:
49 raise ValueError(
50 f"Plugin name is incorrect: expected '{plugin_name}', got '{model.name}'. The author needs to correct this in the plugin.yaml file."
51 )
52 return model
53
54
55 def check_plugin_conflict(name: str) -> None:
56 """Raise ValueError if a plugin with this name already exists in usr/plugins/."""
57 dest = os.path.join(_get_user_plugins_dir(), name)
58 if os.path.exists(dest):
59 raise ValueError(f"Plugin '{name}' is already installed")
60
61
62 def _find_plugin_root(extracted_dir: str) -> str:
63 """Walk extracted directory to find the parent of plugin.yaml.
64 Returns absolute path to the plugin root directory."""
65 for root, dirs, dir_files in os.walk(extracted_dir):
66 if META_FILE_NAME in dir_files:
67 return root
68 raise ValueError(f"No {META_FILE_NAME} found in the uploaded archive")
69
70
71 def install_uploaded_zip(plugin_file: FileStorage) -> dict:
72 """Persist an uploaded ZIP temporarily and install it."""
73 original_filename = Path((plugin_file.filename or "").strip()).name
74 if not original_filename:
75 raise ValueError("No file selected")
76
77 tmp_dir = Path(files.get_abs_path("tmp", "plugin_uploads"))
78 tmp_dir.mkdir(parents=True, exist_ok=True)
79
80 temp_name = secure_filename(original_filename) or "plugin.zip"
81 if not temp_name.lower().endswith(".zip"):
82 temp_name = f"{temp_name}.zip"
83
84 unique = uuid.uuid4().hex[:8]
85 stamp = time.strftime("%Y%m%d_%H%M%S")
86 tmp_path = str(tmp_dir / f"plugin_{stamp}_{unique}_{temp_name}")
87 plugin_file.save(tmp_path)
88
89 return install_from_zip(tmp_path, original_filename=original_filename)
90
91
92 def install_from_zip(zip_path: str, original_filename: str | None = None) -> dict:
93 """Extract ZIP, find plugin.yaml, move its parent to usr/plugins/.
94 Returns dict with plugin name and metadata.
95 Cleans up tmp files regardless of outcome."""
96 temp_name = f"tmp_plugin_{time.strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}"
97 extract_dir = files.get_abs_path(files.TEMP_DIR, "plugin_installs", temp_name)
98 extract_dir = files.create_dir_safe(extract_dir)
99 dest = ""
100
101 try:
102 try:
103 # Extract with path traversal protection
104 with zipfile.ZipFile(zip_path, "r") as z:
105 for member in z.namelist():
106 member_path = os.path.realpath(os.path.join(extract_dir, member))
107 if not (files.is_in_dir(member_path, extract_dir)):
108 raise ValueError(f"Unsafe path in archive: {member}")
109 z.extractall(extract_dir)
110
111 # Find plugin.yaml
112 plugin_root = _find_plugin_root(extract_dir)
113 meta = validate_plugin_dir(plugin_root)
114 plugin_name = _get_plugin_name(meta)
115
116 check_plugin_conflict(plugin_name)
117
118 # Move to usr/plugins/
119 dest = os.path.join(_get_user_plugins_dir(), plugin_name)
120 files.create_dir(os.path.dirname(dest))
121 files.move_dir(plugin_root, dest)
122 except Exception as e:
123 print_style.PrintStyle.error(f"Failed to validate plugin: {e}")
124 files.delete_dir(extract_dir)
125 raise
126
127 # run installation hook
128 try:
129 run_install_hook(plugin_name)
130 except Exception as e:
131 print_style.PrintStyle.error(
132 f"Failed to run installation hook for {plugin_name}: {e}"
133 )
134 files.delete_dir(dest)
135 raise
136
137
138 # does it have python files?
139 python_change = bool(files.find_existing_paths_by_pattern(dest+"/**/*.py"))
140
141 after_plugin_change([plugin_name], python_change=python_change)
142
143 return {
144 "success": True,
145 "plugin_name": plugin_name,
146 "title": meta.title or plugin_name,
147 "path": files.deabsolute_path(dest),
148 }
149 finally:
150 # Cleanup: extracted files and the archive
151 try:
152 files.delete_dir(extract_dir)
153 files.delete_file(zip_path)
154 except Exception as e:
155 pass
156
157
158 def _download_thumbnail(thumbnail_url: str, plugin_dir: str) -> None:
159 """Download thumbnail from URL to plugin_dir/webui/thumbnail.<ext>. Non-fatal."""
160 try:
161 if not thumbnail_url:
162 return
163 from urllib.parse import urlparse
164 parsed = urlparse(thumbnail_url)
165 if parsed.scheme not in ("http", "https"):
166 return
167 _allowed_exts = {"png", "jpg", "jpeg", "gif", "webp"}
168 url_path = parsed.path.lower()
169 ext = url_path.rsplit(".", 1)[-1] if "." in url_path else ""
170 if ext not in _allowed_exts:
171 ext = "png"
172 webui_dir = Path(plugin_dir) / "webui"
173 webui_dir.mkdir(parents=True, exist_ok=True)
174 dest = webui_dir / f"thumbnail.{ext}"
175 req = urllib.request.Request(thumbnail_url, headers={"User-Agent": "AgentZero"})
176 with urllib.request.urlopen(req, timeout=10) as resp:
177 dest.write_bytes(resp.read())
178 except Exception as e:
179 print_style.PrintStyle.warning(f"Failed to download plugin thumbnail: {e}")
180
181
182 def install_from_git(url: str, token: str | None = None, plugin_name: str = "", thumbnail_url: str = "") -> dict:
183 """Clone git repo into usr/plugins/, validate plugin.yaml.
184 Returns dict with plugin name and metadata."""
185 from helpers.git import clone_repo
186
187 temp_name = f"tmp_plugin_{time.strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}"
188 git_dir = files.get_abs_path(files.TEMP_DIR, "plugins_installer", temp_name)
189
190 try:
191 files.create_dir_safe(git_dir)
192 clone_repo(url, git_dir, token=token or None)
193 meta = validate_plugin_dir(git_dir, plugin_name=plugin_name)
194 plugin_name = _get_plugin_name(meta)
195 check_plugin_conflict(plugin_name)
196 final_dir = os.path.join(_get_user_plugins_dir(), plugin_name)
197 files.move_dir(git_dir, final_dir)
198 except Exception as e:
199 # No plugin.yaml — remove cloned repo
200 print_style.PrintStyle.error(f"Failed to validate plugin: {e}")
201 files.delete_dir(git_dir)
202 raise
203
204 _download_thumbnail(thumbnail_url, final_dir)
205
206 # run installation hook
207 try:
208 run_install_hook(plugin_name)
209 except Exception as e:
210 print_style.PrintStyle.error(
211 f"Failed to run installation hook for {plugin_name}: {e}"
212 )
213 files.delete_dir(final_dir)
214 raise
215
216 # does it have python files?
217 python_change = bool(files.find_existing_paths_by_pattern(final_dir+"/**/*.py"))
218
219 after_plugin_change([plugin_name],python_change=python_change)
220
221 return {
222 "success": True,
223 "plugin_name": plugin_name,
224 "title": meta.title or plugin_name,
225 "path": files.deabsolute_path(final_dir),
226 }
227
228
229 def update_from_git(plugin_name: str) -> dict:
230 plugin_name = (plugin_name or "").strip()
231 if not plugin_name:
232 raise ValueError("Missing plugin_name")
233
234 plugin_dir = plugins.find_plugin_dir(plugin_name)
235 if not plugin_dir:
236 raise ValueError("Plugin not found")
237
238 custom_plugins_dir = files.get_abs_path(files.USER_DIR, files.PLUGINS_DIR)
239 if not files.is_in_dir(plugin_dir, custom_plugins_dir):
240 raise ValueError("Only custom plugins can be updated")
241
242 try:
243 run_pre_update_hook(plugin_name)
244 except Exception as e:
245 print_style.PrintStyle.error(
246 f"Failed to run pre-update hook for {plugin_name}: {e}"
247 )
248 raise
249
250 try:
251 repo = git.update_repo(plugin_dir)
252 meta = plugins.get_plugin_meta(plugin_name)
253 except git.DirtyTreeConflictError as e:
254 print_style.PrintStyle.error(f"Failed to update plugin: {e}")
255 return {
256 "ok": False,
257 "success": False,
258 "error": str(e),
259 "error_kind": "dirty_tree_conflict",
260 "plugin_name": plugin_name,
261 "conflicting_files": e.conflicting_files,
262 }
263 except Exception as e:
264 print_style.PrintStyle.error(f"Failed to update plugin: {e}")
265 raise
266
267 try:
268 run_install_hook(plugin_name)
269 except Exception as e:
270 print_style.PrintStyle.error(
271 f"Failed to run installation hook for {plugin_name}: {e}"
272 )
273 raise
274
275 after_plugin_change([plugin_name])
276 head = repo.head.commit
277
278 return {
279 "ok": True,
280 "success": True,
281 "plugin_name": plugin_name,
282 "title": meta.title if meta else plugin_name,
283 "path": files.deabsolute_path(plugin_dir),
284 "current_commit": head.hexsha,
285 "current_commit_timestamp": datetime.fromtimestamp(
286 head.committed_date,
287 tz=Localization.get().get_tzinfo(),
288 ).strftime("%Y-%m-%d %H:%M:%S %Z"),
289 "version": getattr(meta, "version", "") or "",
290 "branch": repo.active_branch.name if not repo.head.is_detached else "",
291 "remote_url": git.strip_auth_from_url(repo.remotes.origin.url) if repo.remotes else "",
292 "directory_name": Path(plugin_dir).name,
293 }
294
295
296 def run_install_hook(plugin_name: str):
297 return plugins.call_plugin_hook(plugin_name, "install")
298
299 def run_pre_update_hook(plugin_name: str):
300 return plugins.call_plugin_hook(plugin_name, "pre_update")
301
302 def get_plugin_hub_index(force: bool = False) -> dict[str, Any]:
303 """Return the plugin index plus installed Plugin Hub keys."""
304 index_data = fetch_plugin_index(force=force)
305 if not isinstance(index_data, dict):
306 raise ValueError("Plugin index response was not a JSON object")
307
308 plugins = index_data.get("plugins")
309 if not isinstance(plugins, dict):
310 raise ValueError("Plugin index payload is missing a valid 'plugins' map")
311
312 from helpers.plugins import find_plugin_dir
313
314 installed_dirs = set(get_plugins_list())
315 installed_keys: list[str] = []
316 _thumb_exts = ("png", "jpg", "jpeg", "gif", "webp")
317
318 for key, plugin_data in plugins.items():
319 if not isinstance(plugin_data, dict):
320 continue
321 if key not in installed_dirs:
322 continue
323 installed_keys.append(key)
324
325 # Backfill thumbnail for plugins installed before this feature existed
326 plugin_dir = find_plugin_dir(key)
327 if not plugin_dir:
328 continue
329 webui_dir = Path(plugin_dir) / "webui"
330 has_thumb = any((webui_dir / f"thumbnail.{ext}").is_file() for ext in _thumb_exts)
331 if has_thumb:
332 continue
333 thumb_url = plugin_data.get("thumbnail") or ""
334 if not thumb_url:
335 from urllib.parse import urlparse
336 raw_base = None
337 github = plugin_data.get("github") or ""
338 if github:
339 try:
340 parsed = urlparse(github)
341 if parsed.netloc == "github.com":
342 parts = parsed.path.strip("/").split("/")
343 if len(parts) >= 2:
344 raw_base = f"https://raw.githubusercontent.com/{parts[0]}/{parts[1]}"
345 except Exception:
346 pass
347 if raw_base:
348 thumb_url = f"{raw_base}/main/thumbnail.png"
349 if thumb_url:
350 _download_thumbnail(thumb_url, plugin_dir)
351
352 return {"index": index_data, "installed_plugins": installed_keys}
353
354
355 def fetch_plugin_index(force: bool = False) -> dict:
356 """Download the plugin index from GitHub releases."""
357 index_url = "https://github.com/agent0ai/a0-plugins/releases/download/generated-index/index.json"
358 if force:
359 separator = "&" if "?" in index_url else "?"
360 index_url = f"{index_url}{separator}ts={time.time_ns()}"
361
362 headers = {"User-Agent": "AgentZero"}
363 if force:
364 headers["Cache-Control"] = "no-cache"
365 headers["Pragma"] = "no-cache"
366
367 req = urllib.request.Request(index_url, headers=headers)
368 with urllib.request.urlopen(req, timeout=30) as resp:
369 data = json.loads(resp.read().decode())
370 return data