from __future__ import annotations import base64 import re from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path import aiofiles import yaml from pydantic import ValidationError from mcp_docugen.models import TemplateAsset, TemplateFrontmatter, TemplateSummary _SLUG_RE = re.compile(r"^[a-z0-9][a-z0-9-]*$") _ASSET_FILENAME_RE = re.compile(r"^[A-Za-z0-9._-]+$") _FRONTMATTER_DELIM = "---" class TemplateNotFound(Exception): pass class TemplateAlreadyExists(Exception): pass class InvalidFrontmatter(Exception): pass class InvalidTemplateName(ValueError): pass @dataclass class LoadedTemplate: frontmatter: TemplateFrontmatter body: str assets: list[TemplateAsset] updated_at: datetime class TemplateStore: def __init__(self, base_dir: Path) -> None: self.base_dir = Path(base_dir) self.base_dir.mkdir(parents=True, exist_ok=True) async def create( self, name: str, frontmatter: TemplateFrontmatter, body: str, assets: list[dict] | None = None, ) -> None: self._validate_name(name) tdir = self.base_dir / name if tdir.exists(): raise TemplateAlreadyExists(name) tdir.mkdir(parents=True) (tdir / "assets").mkdir() if assets: self._validate_assets(assets) await self._write_template_file(tdir, frontmatter, body) if assets: await self._write_assets(tdir, assets) async def update( self, name: str, frontmatter: TemplateFrontmatter, body: str, assets: list[dict] | None = None, ) -> None: self._validate_name(name) tdir = self.base_dir / name if not tdir.exists(): raise TemplateNotFound(name) await self._write_template_file(tdir, frontmatter, body) if assets is not None: assets_dir = tdir / "assets" for f in assets_dir.iterdir(): f.unlink() await self._write_assets(tdir, assets) async def delete(self, name: str) -> None: self._validate_name(name) tdir = self.base_dir / name if not tdir.exists(): raise TemplateNotFound(name) for f in (tdir / "assets").iterdir(): f.unlink() (tdir / "assets").rmdir() (tdir / "template.md").unlink() tdir.rmdir() async def get(self, name: str) -> LoadedTemplate: self._validate_name(name) tdir = self.base_dir / name tpath = tdir / "template.md" if not tpath.exists(): raise TemplateNotFound(name) frontmatter, body = await self._read_template_file(tpath) assets = self._list_assets(tdir) stat = tpath.stat() updated_at = datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc) return LoadedTemplate( frontmatter=frontmatter, body=body, assets=assets, updated_at=updated_at ) async def list(self) -> list[TemplateSummary]: summaries = [] for tdir in sorted(self.base_dir.iterdir()): if not tdir.is_dir(): continue tpath = tdir / "template.md" if not tpath.exists(): continue try: frontmatter, _ = await self._read_template_file(tpath) except InvalidFrontmatter: continue stat = tpath.stat() summaries.append( TemplateSummary( name=frontmatter.name, description=frontmatter.description, updated_at=datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc), ) ) return summaries async def read_asset(self, template_name: str, filename: str) -> bytes: self._validate_name(template_name) self._validate_asset_filename(filename) path = self.base_dir / template_name / "assets" / filename if not path.exists(): raise FileNotFoundError(filename) async with aiofiles.open(path, "rb") as f: return await f.read() def asset_path(self, template_name: str, filename: str) -> Path: self._validate_name(template_name) self._validate_asset_filename(filename) return self.base_dir / template_name / "assets" / filename def _validate_name(self, name: str) -> None: if not _SLUG_RE.match(name): raise InvalidTemplateName(name) def _validate_asset_filename(self, filename: str) -> None: if not _ASSET_FILENAME_RE.match(filename): raise ValueError(f"invalid asset filename: {filename!r}") def _validate_assets(self, assets: list[dict]) -> None: for asset in assets: self._validate_asset_filename(asset["filename"]) async def _write_template_file( self, tdir: Path, frontmatter: TemplateFrontmatter, body: str ) -> None: yaml_text = yaml.safe_dump( frontmatter.model_dump(exclude_none=True), sort_keys=False ) content = f"{_FRONTMATTER_DELIM}\n{yaml_text}{_FRONTMATTER_DELIM}\n{body}" async with aiofiles.open(tdir / "template.md", "w") as f: await f.write(content) async def _read_template_file( self, path: Path ) -> tuple[TemplateFrontmatter, str]: async with aiofiles.open(path) as f: raw = await f.read() if not raw.startswith(_FRONTMATTER_DELIM): raise InvalidFrontmatter("missing opening '---'") parts = raw.split(_FRONTMATTER_DELIM, 2) if len(parts) < 3: raise InvalidFrontmatter("missing closing '---'") yaml_text = parts[1] body = parts[2].lstrip("\n") try: data = yaml.safe_load(yaml_text) or {} frontmatter = TemplateFrontmatter(**data) except (yaml.YAMLError, ValidationError) as exc: raise InvalidFrontmatter(str(exc)) from exc return frontmatter, body async def _write_assets(self, tdir: Path, assets: list[dict]) -> None: for asset in assets: filename = asset["filename"] self._validate_asset_filename(filename) data = base64.b64decode(asset["data_b64"]) async with aiofiles.open(tdir / "assets" / filename, "wb") as f: await f.write(data) def _list_assets(self, tdir: Path) -> list[TemplateAsset]: assets_dir = tdir / "assets" if not assets_dir.exists(): return [] out = [] for f in sorted(assets_dir.iterdir()): if not f.is_file(): continue out.append( TemplateAsset( filename=f.name, mime=_guess_mime(f.name), size_bytes=f.stat().st_size, ) ) return out def _guess_mime(filename: str) -> str: ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else "" return { "png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg", "webp": "image/webp", }.get(ext, "application/octet-stream")