| 1 | """Config-driven local and S3-compatible file storage.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import base64 |
| 6 | import mimetypes |
| 7 | import re |
| 8 | import shutil |
| 9 | from hashlib import sha256 |
| 10 | from pathlib import Path |
| 11 | from typing import Any |
| 12 | from urllib.parse import quote, unquote, urlparse |
| 13 | |
| 14 | |
| 15 | def _safe_segment(value: str, *, fallback: str) -> str: |
| 16 | cleaned = re.sub(r"[^A-Za-z0-9._-]+", "_", value.strip()).strip("._") |
| 17 | return cleaned or fallback |
| 18 | |
| 19 | |
| 20 | def _local_root(config: Any, workspace: Path) -> Path: |
| 21 | configured = Path(str(config.directory).strip()).expanduser() |
| 22 | return (configured if configured.is_absolute() else workspace / configured).resolve() |
| 23 | |
| 24 | |
| 25 | class LocalFilePublisher: |
| 26 | """Copy files into the configured local store and return gateway URLs.""" |
| 27 | |
| 28 | def __init__(self, config: Any, *, workspace: Path, work_id: str) -> None: |
| 29 | self.base_url = str(config.base_url).strip().rstrip("/") |
| 30 | self.route_prefix = "/" + str(config.route_prefix).strip().strip("/") |
| 31 | self.work_id = _safe_segment(work_id, fallback="work") |
| 32 | self.root = _local_root(config, workspace.expanduser().resolve()) |
| 33 | self.delete_local_after_upload = False |
| 34 | |
| 35 | def _relative_path(self, name: str, *, digest: str) -> Path: |
| 36 | parts = [ |
| 37 | _safe_segment(part, fallback="asset") |
| 38 | for part in name.replace("\\", "/").split("/") |
| 39 | if part.strip() |
| 40 | ] |
| 41 | relative = Path(*parts) if parts else Path("asset") |
| 42 | return Path(self.work_id) / relative.with_name( |
| 43 | f"{relative.stem}_{digest}{relative.suffix}" |
| 44 | ) |
| 45 | |
| 46 | def __call__(self, local: str, name: str) -> str: |
| 47 | source = Path(local).expanduser().resolve() |
| 48 | if not source.is_file() or source.stat().st_size <= 0: |
| 49 | raise RuntimeError(f"file is missing or empty: {source}") |
| 50 | digest = sha256(source.read_bytes()).hexdigest()[:12] |
| 51 | relative = self._relative_path(name, digest=digest) |
| 52 | target = self.root / relative |
| 53 | target.parent.mkdir(parents=True, exist_ok=True) |
| 54 | if source != target.resolve(strict=False): |
| 55 | shutil.copy2(source, target) |
| 56 | path = quote(relative.as_posix(), safe="/-_.~") |
| 57 | url = f"{self.route_prefix}/{path}" |
| 58 | return f"{self.base_url}{url}" if self.base_url else url |
| 59 | |
| 60 | |
| 61 | class S3FilePublisher: |
| 62 | """Upload files to explicitly configured S3-compatible storage.""" |
| 63 | |
| 64 | def __init__(self, config: Any, *, work_id: str, client: Any | None = None) -> None: |
| 65 | self.bucket = str(config.bucket).strip() |
| 66 | self.region = str(config.region).strip() |
| 67 | self.endpoint_url = str(config.endpoint_url).strip().rstrip("/") |
| 68 | self.public_base_url = str(config.public_base_url).strip().rstrip("/") |
| 69 | self.key_prefix = str(config.key_prefix).strip().strip("/") |
| 70 | self.work_id = _safe_segment(work_id, fallback="work") |
| 71 | self._config = config |
| 72 | self._client = client |
| 73 | access_key = str(config.access_key_id).strip() |
| 74 | secret_key = str(config.secret_access_key).strip() |
| 75 | if not all( |
| 76 | (self.bucket, self.endpoint_url, self.public_base_url, access_key, secret_key) |
| 77 | ): |
| 78 | raise RuntimeError( |
| 79 | "tools.fileStorage.outbound.s3 requires endpointUrl, publicBaseUrl, bucket, " |
| 80 | "accessKeyId, and secretAccessKey" |
| 81 | ) |
| 82 | |
| 83 | def _client_for_upload(self) -> Any: |
| 84 | if self._client is None: |
| 85 | try: |
| 86 | import boto3 |
| 87 | from botocore.config import Config as BotoConfig |
| 88 | except ImportError as exc: |
| 89 | raise RuntimeError( |
| 90 | "S3-compatible file upload requires boto3; run setup_local.sh again" |
| 91 | ) from exc |
| 92 | kwargs: dict[str, Any] = { |
| 93 | "endpoint_url": self.endpoint_url, |
| 94 | "region_name": self.region or None, |
| 95 | "aws_access_key_id": str(self._config.access_key_id).strip(), |
| 96 | "aws_secret_access_key": str(self._config.secret_access_key).strip(), |
| 97 | "config": BotoConfig( |
| 98 | signature_version="s3v4", |
| 99 | s3={"addressing_style": str(self._config.addressing_style)}, |
| 100 | request_checksum_calculation="when_required", |
| 101 | response_checksum_validation="when_required", |
| 102 | ), |
| 103 | } |
| 104 | session_token = str(self._config.session_token).strip() |
| 105 | if session_token: |
| 106 | kwargs["aws_session_token"] = session_token |
| 107 | self._client = boto3.client("s3", **kwargs) |
| 108 | return self._client |
| 109 | |
| 110 | def _object_key(self, name: str, *, digest: str) -> str: |
| 111 | parts = [ |
| 112 | _safe_segment(part, fallback="asset") |
| 113 | for part in name.replace("\\", "/").split("/") |
| 114 | if part.strip() |
| 115 | ] |
| 116 | relative = "/".join(parts) or "asset" |
| 117 | path = Path(relative) |
| 118 | relative = str(path.with_name(f"{path.stem}_{digest}{path.suffix}")) |
| 119 | return "/".join( |
| 120 | part for part in (self.key_prefix, self.work_id, relative) if part |
| 121 | ) |
| 122 | |
| 123 | def __call__(self, local: str, name: str) -> str: |
| 124 | source = Path(local).expanduser().resolve() |
| 125 | if not source.is_file() or source.stat().st_size <= 0: |
| 126 | raise RuntimeError(f"file is missing or empty: {source}") |
| 127 | digest = sha256(source.read_bytes()).hexdigest()[:12] |
| 128 | object_key = self._object_key(name, digest=digest) |
| 129 | content_type = mimetypes.guess_type(source.name)[0] or "application/octet-stream" |
| 130 | self._client_for_upload().upload_file( |
| 131 | str(source), |
| 132 | self.bucket, |
| 133 | object_key, |
| 134 | ExtraArgs={"ContentType": content_type}, |
| 135 | ) |
| 136 | return f"{self.public_base_url}/{quote(object_key, safe='/-_.~')}" |
| 137 | |
| 138 | |
| 139 | def configured_file_publisher( |
| 140 | work_id: str, |
| 141 | *, |
| 142 | storage: Any | None = None, |
| 143 | workspace: Path | None = None, |
| 144 | ) -> Any: |
| 145 | """Build the canonical local file publisher.""" |
| 146 | if storage is None or workspace is None: |
| 147 | from nanobot.config.loader import load_config |
| 148 | |
| 149 | config = load_config() |
| 150 | storage = storage or config.tools.file_storage |
| 151 | workspace = workspace or Path(config.agents.defaults.workspace) |
| 152 | return LocalFilePublisher( |
| 153 | storage.local, |
| 154 | workspace=workspace, |
| 155 | work_id=work_id, |
| 156 | ) |
| 157 | |
| 158 | |
| 159 | def resolve_local_asset_path(url: str, *, workspace: Path, config: Any) -> Path | None: |
| 160 | """Resolve one configured local-store URL without allowing traversal.""" |
| 161 | parsed = urlparse(url) |
| 162 | path = unquote(parsed.path or url) |
| 163 | prefix = "/" + str(config.route_prefix).strip().strip("/") + "/" |
| 164 | if not path.startswith(prefix): |
| 165 | return None |
| 166 | relative = path[len(prefix):] |
| 167 | root = _local_root(config, workspace.expanduser().resolve()) |
| 168 | candidate = (root / relative).resolve() |
| 169 | try: |
| 170 | candidate.relative_to(root) |
| 171 | except ValueError: |
| 172 | return None |
| 173 | return candidate if candidate.is_file() else None |
| 174 | |
| 175 | |
| 176 | def local_asset_data_uri(url: str, *, workspace: Path, config: Any) -> str: |
| 177 | """Convert a configured local-store URL to an inline payload.""" |
| 178 | path = resolve_local_asset_path(url, workspace=workspace, config=config) |
| 179 | if path is None: |
| 180 | return url |
| 181 | mime = mimetypes.guess_type(path.name)[0] or "application/octet-stream" |
| 182 | encoded = base64.b64encode(path.read_bytes()).decode("ascii") |
| 183 | return f"data:{mime};base64,{encoded}" |
| 184 | |
| 185 | |
| 186 | def outbound_file_url( |
| 187 | url: str, |
| 188 | *, |
| 189 | workspace: Path, |
| 190 | work_id: str, |
| 191 | name: str, |
| 192 | storage: Any | None = None, |
| 193 | ) -> str: |
| 194 | """Prepare a locally stored file for an outbound service request.""" |
| 195 | if storage is None: |
| 196 | from nanobot.config.loader import load_config |
| 197 | |
| 198 | storage = load_config().tools.file_storage |
| 199 | path = resolve_local_asset_path(url, workspace=workspace, config=storage.local) |
| 200 | if path is None: |
| 201 | return url |
| 202 | if storage.outbound.backend == "s3": |
| 203 | return S3FilePublisher(storage.outbound.s3, work_id=work_id)(str(path), name) |
| 204 | return local_asset_data_uri(url, workspace=workspace, config=storage.local) |
| 205 |