| 1 | from __future__ import annotations |
| 2 | |
| 3 | import os |
| 4 | import shutil |
| 5 | import stat |
| 6 | import tempfile |
| 7 | import time |
| 8 | import zipfile |
| 9 | from dataclasses import dataclass |
| 10 | from pathlib import Path |
| 11 | from typing import Iterable, List, Literal, Optional, Tuple |
| 12 | |
| 13 | from helpers import files |
| 14 | from helpers.skills import discover_skill_md_files |
| 15 | |
| 16 | |
| 17 | ConflictPolicy = Literal["skip", "overwrite", "rename"] |
| 18 | |
| 19 | # Project skills folder name (inside .a0proj) |
| 20 | PROJECT_SKILLS_DIR = "skills" |
| 21 | |
| 22 | |
| 23 | @dataclass(slots=True) |
| 24 | class ImportPlanItem: |
| 25 | src_root: Path |
| 26 | src_skill_dir: Path |
| 27 | dest_skill_dir: Path |
| 28 | |
| 29 | |
| 30 | @dataclass(slots=True) |
| 31 | class ImportResult: |
| 32 | imported: List[Path] |
| 33 | skipped: List[Path] |
| 34 | source_root: Path |
| 35 | destination_root: Path |
| 36 | namespace: str |
| 37 | |
| 38 | |
| 39 | def _is_within(child: Path, parent: Path) -> bool: |
| 40 | try: |
| 41 | child.resolve().relative_to(parent.resolve()) |
| 42 | return True |
| 43 | except Exception: |
| 44 | return False |
| 45 | |
| 46 | |
| 47 | def _derive_namespace(source: Path) -> str: |
| 48 | # Use stem for zip, name for directory |
| 49 | return (source.stem or source.name or "import").strip() |
| 50 | |
| 51 | |
| 52 | def _candidate_skill_roots(source_dir: Path) -> List[Path]: |
| 53 | """ |
| 54 | Heuristics to find likely skill roots inside a repo/pack: |
| 55 | - <source>/skills |
| 56 | - <source>/plugins/*/skills (Claude Code style) |
| 57 | - fallback: <source> |
| 58 | """ |
| 59 | candidates: List[Path] = [] |
| 60 | |
| 61 | direct = source_dir / "skills" |
| 62 | if direct.is_dir() and discover_skill_md_files(direct): |
| 63 | candidates.append(direct) |
| 64 | |
| 65 | plugins = source_dir / "plugins" |
| 66 | if plugins.is_dir(): |
| 67 | for child in plugins.iterdir(): |
| 68 | if not child.is_dir(): |
| 69 | continue |
| 70 | skills_dir = child / "skills" |
| 71 | if skills_dir.is_dir() and discover_skill_md_files(skills_dir): |
| 72 | candidates.append(skills_dir) |
| 73 | |
| 74 | # Deduplicate while preserving order |
| 75 | unique: List[Path] = [] |
| 76 | seen = set() |
| 77 | for c in candidates: |
| 78 | key = str(c.resolve()) |
| 79 | if key not in seen: |
| 80 | seen.add(key) |
| 81 | unique.append(c) |
| 82 | |
| 83 | return unique or [source_dir] |
| 84 | |
| 85 | |
| 86 | def _safe_extract_zip(zip_path: Path, target: Path) -> None: |
| 87 | with zipfile.ZipFile(zip_path, "r") as archive: |
| 88 | for member in archive.infolist(): |
| 89 | name = member.filename |
| 90 | if not name or name.startswith(("/", "\\")) or "\\" in name: |
| 91 | raise ValueError(f"Unsafe zip entry path: {name!r}") |
| 92 | |
| 93 | mode = member.external_attr >> 16 |
| 94 | if stat.S_ISLNK(mode): |
| 95 | raise ValueError(f"Refusing to extract symlink from zip: {name!r}") |
| 96 | |
| 97 | destination = target / name |
| 98 | if not _is_within(destination, target): |
| 99 | raise ValueError(f"Unsafe zip entry path: {name!r}") |
| 100 | |
| 101 | archive.extractall(target) |
| 102 | |
| 103 | |
| 104 | def extract_skills_zip( |
| 105 | zip_path: Path, |
| 106 | *, |
| 107 | tmp_subdir: str = "skill_imports", |
| 108 | prefix: str = "import", |
| 109 | ) -> tuple[Path, Path]: |
| 110 | """ |
| 111 | Extract a zip into a temp folder inside Agent Zero's tmp directory. |
| 112 | Returns (source_root, cleanup_root). |
| 113 | """ |
| 114 | base_tmp = Path(files.get_abs_path("tmp", tmp_subdir)) |
| 115 | base_tmp.mkdir(parents=True, exist_ok=True) |
| 116 | stamp = time.strftime("%Y%m%d_%H%M%S") |
| 117 | target = base_tmp / f"{prefix}_{zip_path.stem}_{stamp}" |
| 118 | target.mkdir(parents=True, exist_ok=True) |
| 119 | |
| 120 | try: |
| 121 | _safe_extract_zip(zip_path, target) |
| 122 | except Exception: |
| 123 | shutil.rmtree(target, ignore_errors=True) |
| 124 | raise |
| 125 | |
| 126 | # If zip contains a single top-level folder, treat that as the root |
| 127 | children = [p for p in target.iterdir()] |
| 128 | if len(children) == 1 and children[0].is_dir(): |
| 129 | return children[0], target |
| 130 | return target, target |
| 131 | |
| 132 | |
| 133 | def _unzip_to_temp_dir(zip_path: Path) -> Path: |
| 134 | """ |
| 135 | Extract a zip into a temp folder under tmp/skill_imports (inside Agent Zero base dir). |
| 136 | Returns the extraction root folder. |
| 137 | """ |
| 138 | source_root, _cleanup_root = extract_skills_zip(zip_path) |
| 139 | return source_root |
| 140 | |
| 141 | |
| 142 | def build_import_plan( |
| 143 | source: Path, |
| 144 | dest_root: Path, |
| 145 | *, |
| 146 | namespace: Optional[str] = None, |
| 147 | ) -> Tuple[List[ImportPlanItem], Path]: |
| 148 | """ |
| 149 | Build a copy plan for importing skills from a source folder. |
| 150 | |
| 151 | Returns: (plan_items, source_root_dir_used_for_scan) |
| 152 | """ |
| 153 | source_dir = source |
| 154 | roots = _candidate_skill_roots(source_dir) |
| 155 | plan: List[ImportPlanItem] = [] |
| 156 | ns = (namespace or _derive_namespace(source)).strip() |
| 157 | dest_ns_root = dest_root / ns |
| 158 | |
| 159 | for root in roots: |
| 160 | for skill_md in discover_skill_md_files(root): |
| 161 | skill_dir = skill_md.parent |
| 162 | # Skip if the skill dir is already inside destination (prevents recursive import) |
| 163 | if _is_within(skill_dir, dest_root): |
| 164 | continue |
| 165 | try: |
| 166 | rel = skill_dir.resolve().relative_to(root.resolve()) |
| 167 | except Exception: |
| 168 | # If relative fails due to symlink oddities, just use leaf folder name |
| 169 | rel = Path(skill_dir.name) |
| 170 | dest_dir = dest_ns_root / rel |
| 171 | plan.append(ImportPlanItem(src_root=root, src_skill_dir=skill_dir, dest_skill_dir=dest_dir)) |
| 172 | |
| 173 | # Deduplicate by destination path (keep first occurrence) |
| 174 | seen_dest = set() |
| 175 | deduped: List[ImportPlanItem] = [] |
| 176 | for item in plan: |
| 177 | key = str(item.dest_skill_dir.resolve()) |
| 178 | if key in seen_dest: |
| 179 | continue |
| 180 | seen_dest.add(key) |
| 181 | deduped.append(item) |
| 182 | |
| 183 | return deduped, roots[0] |
| 184 | |
| 185 | |
| 186 | def _resolve_conflict(dest: Path, policy: ConflictPolicy) -> Tuple[Path, bool]: |
| 187 | """ |
| 188 | Returns (final_dest_path, should_copy). |
| 189 | """ |
| 190 | if not dest.exists(): |
| 191 | return dest, True |
| 192 | |
| 193 | if policy == "skip": |
| 194 | return dest, False |
| 195 | |
| 196 | if policy == "overwrite": |
| 197 | shutil.rmtree(dest) |
| 198 | return dest, True |
| 199 | |
| 200 | # rename |
| 201 | i = 2 |
| 202 | while True: |
| 203 | candidate = dest.with_name(f"{dest.name}_{i}") |
| 204 | if not candidate.exists(): |
| 205 | return candidate, True |
| 206 | i += 1 |
| 207 | |
| 208 | |
| 209 | def get_project_skills_folder(project_name: str) -> Path: |
| 210 | """Get the skills folder path for a project.""" |
| 211 | from helpers.projects import get_project_meta |
| 212 | return Path(get_project_meta(project_name, PROJECT_SKILLS_DIR)) |
| 213 | |
| 214 | |
| 215 | def get_agent_profile_skills_folder(profile_name: str) -> Path: |
| 216 | return Path(files.get_abs_path("usr", "agents", profile_name, "skills")) |
| 217 | |
| 218 | |
| 219 | def get_project_agent_profile_skills_folder(project_name: str, profile_name: str) -> Path: |
| 220 | from helpers.projects import get_project_meta |
| 221 | return Path(get_project_meta(project_name, "agents", profile_name, "skills")) |
| 222 | |
| 223 | |
| 224 | def resolve_skills_destination_root( |
| 225 | project_name: Optional[str], |
| 226 | agent_profile: Optional[str], |
| 227 | ) -> Path: |
| 228 | if project_name and agent_profile: |
| 229 | return get_project_agent_profile_skills_folder(project_name, agent_profile) |
| 230 | if project_name: |
| 231 | return get_project_skills_folder(project_name) |
| 232 | if agent_profile: |
| 233 | return get_agent_profile_skills_folder(agent_profile) |
| 234 | return Path(files.get_abs_path("usr", "skills")) |
| 235 | |
| 236 | |
| 237 | def import_skills( |
| 238 | source_path: str, |
| 239 | *, |
| 240 | namespace: Optional[str] = None, |
| 241 | conflict: ConflictPolicy = "skip", |
| 242 | dry_run: bool = False, |
| 243 | project_name: Optional[str] = None, |
| 244 | agent_profile: Optional[str] = None, |
| 245 | ) -> ImportResult: |
| 246 | """ |
| 247 | Import external Skills into usr/skills/<namespace>/... |
| 248 | |
| 249 | - source_path can be a directory or a .zip file |
| 250 | - Uses heuristics to detect the Skills root(s) |
| 251 | - Copies each skill folder (parent of SKILL.md) as-is |
| 252 | """ |
| 253 | src = Path(source_path).expanduser() |
| 254 | if not src.is_absolute(): |
| 255 | src = (Path.cwd() / src).resolve() |
| 256 | |
| 257 | if not src.exists(): |
| 258 | raise FileNotFoundError(f"Source not found: {src}") |
| 259 | |
| 260 | dest_root = resolve_skills_destination_root(project_name, agent_profile) |
| 261 | dest_root.mkdir(parents=True, exist_ok=True) |
| 262 | |
| 263 | extracted_root: Optional[Path] = None |
| 264 | source_dir: Path |
| 265 | if src.is_file() and src.suffix.lower() == ".zip": |
| 266 | extracted_root = _unzip_to_temp_dir(src) |
| 267 | source_dir = extracted_root |
| 268 | elif src.is_dir(): |
| 269 | source_dir = src |
| 270 | else: |
| 271 | raise ValueError("Source must be a directory or a .zip file") |
| 272 | |
| 273 | ns = (namespace or _derive_namespace(src)).strip() |
| 274 | if not ns: |
| 275 | ns = "import" |
| 276 | |
| 277 | plan, root_used = build_import_plan(source_dir, dest_root, namespace=ns) |
| 278 | imported: List[Path] = [] |
| 279 | skipped: List[Path] = [] |
| 280 | |
| 281 | for item in plan: |
| 282 | final_dest, should_copy = _resolve_conflict(item.dest_skill_dir, conflict) |
| 283 | if not should_copy: |
| 284 | skipped.append(item.dest_skill_dir) |
| 285 | continue |
| 286 | if dry_run: |
| 287 | imported.append(final_dest) |
| 288 | continue |
| 289 | final_dest.parent.mkdir(parents=True, exist_ok=True) |
| 290 | shutil.copytree(item.src_skill_dir, final_dest) |
| 291 | imported.append(final_dest) |
| 292 | |
| 293 | return ImportResult( |
| 294 | imported=imported, |
| 295 | skipped=skipped, |
| 296 | source_root=root_used, |
| 297 | destination_root=dest_root, |
| 298 | namespace=ns, |
| 299 | ) |