iron/scripts/real_openlist_acceptance.py

from __future__ import annotations

import base64
import json
import os
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass


def main() -> int:
    config = AcceptanceConfig.from_env()
    client = IronClient(config)
    file_class = os.getenv("IRON_ACCEPTANCE_FILE_CLASS", "other")

    print(f"[1/7] logging into Iron at {config.iron_base_url}")
    client.login()

    print(f"[2/7] ensuring backend {config.backend_name!r}")
    backend = client.ensure_webdav_backend()
    print(f"       backend id: {backend['id']}")

    print("[3/7] running backend health check")
    check = client.check_backend(backend["id"])
    if check["status"] != "healthy":
        raise RuntimeError(f"backend health check failed: {json.dumps(check, ensure_ascii=False)}")
    print(f"       healthy: {check['detail']}")

    original_policy = next(policy for policy in client.list_policies() if policy["file_class"] == file_class)
    client.upsert_policy(
        file_class,
        require_local=True,
        stable_replica_count=1,
        opportunistic_replica_count=0,
        preferred_backend_ids=[backend["id"]],
        excluded_backend_ids=[],
        max_non_local_size_bytes=None,
    )

    try:
        print("[4/7] uploading acceptance file through Iron")
        file_id = client.upload_text_file()
        print(f"       file id: {file_id}")

        print("[5/7] waiting for local + OpenList replicas to become ready")
        detail = client.wait_for_ready_replicas(file_id, backend["id"])
        print(f"       replicas ready: {detail['ready_replica_backend_ids']}")

        print("[6/7] verifying download through Iron")
        downloaded = client.download_file(file_id)
        if downloaded != config.test_content.encode("utf-8"):
            raise RuntimeError("downloaded bytes do not match uploaded content")
        print(f"       download ok: {len(downloaded)} bytes")

        print("[7/7] verifying direct WebDAV read from OpenList")
        webdav_replica = next(
            replica for replica in detail["replicas"] if replica["backend_id"] == backend["id"] and replica["status"] == "ready"
        )
        dav_bytes = read_openlist_webdav(
            endpoint_url=config.openlist_dav_url,
            username=config.openlist_username,
            password=config.openlist_password,
            relative_path=f"{config.openlist_root_path}/{webdav_replica['storage_key']}",
        )
        if dav_bytes != config.test_content.encode("utf-8"):
            raise RuntimeError("OpenList WebDAV bytes do not match uploaded content")
        print(f"       webdav ok: {webdav_replica['storage_key']}")

        print("Acceptance passed.")
        return 0
    finally:
        client.upsert_policy(
            file_class,
            require_local=bool(original_policy["require_local"]),
            stable_replica_count=int(original_policy["stable_replica_count"]),
            opportunistic_replica_count=int(original_policy["opportunistic_replica_count"]),
            preferred_backend_ids=list(original_policy["preferred_backend_ids"]),
            excluded_backend_ids=list(original_policy["excluded_backend_ids"]),
            max_non_local_size_bytes=original_policy["max_non_local_size_bytes"],
        )


@dataclass
class AcceptanceConfig:
    iron_base_url: str
    iron_username: str
    iron_password: str
    openlist_dav_url: str
    openlist_username: str
    openlist_password: str
    openlist_root_path: str
    backend_name: str
    directory_id: str
    test_filename: str
    test_content: str
    reconcile_timeout_seconds: float
    reconcile_poll_interval_seconds: float
    http_timeout_seconds: float

    @classmethod
    def from_env(cls) -> "AcceptanceConfig":
        default_filename = f"real-openlist-acceptance-{int(time.time())}.txt"
        return cls(
            iron_base_url=os.getenv("IRON_BASE_URL", "http://127.0.0.1:8000").rstrip("/"),
            iron_username=os.getenv("IRON_USERNAME", "admin"),
            iron_password=os.getenv("IRON_PASSWORD", "changeme-iron"),
            openlist_dav_url=os.getenv("OPENLIST_DAV_URL", "http://127.0.0.1:5244/dav").rstrip("/"),
            openlist_username=os.getenv("OPENLIST_USERNAME", "iron"),
            openlist_password=os.getenv("OPENLIST_PASSWORD", "iron"),
            openlist_root_path=os.getenv("OPENLIST_ROOT_PATH", "aliyun/iron").strip("/"),
            backend_name=os.getenv("IRON_OPENLIST_BACKEND_NAME", "openlist-aliyun"),
            directory_id=os.getenv("IRON_TEST_DIRECTORY_ID", "dir_root"),
            test_filename=os.getenv("IRON_TEST_FILENAME", default_filename),
            test_content=os.getenv("IRON_TEST_CONTENT", "real-openlist-acceptance\n"),
            reconcile_timeout_seconds=float(os.getenv("IRON_RECONCILE_TIMEOUT_SECONDS", "45")),
            reconcile_poll_interval_seconds=float(os.getenv("IRON_RECONCILE_POLL_INTERVAL_SECONDS", "2")),
            http_timeout_seconds=float(os.getenv("IRON_HTTP_TIMEOUT_SECONDS", "30")),
        )


class IronClient:
    def __init__(self, config: AcceptanceConfig) -> None:
        self.config = config
        self.token: str | None = None

    def login(self) -> None:
        response = self._request(
            "POST",
            "/api/auth/login",
            json_body={"username": self.config.iron_username, "password": self.config.iron_password},
        )
        self.token = response["access_token"]

    def ensure_webdav_backend(self) -> dict[str, object]:
        items = self.list_backends()
        for backend in items:
            if backend["name"] == self.config.backend_name:
                return backend
        payload = {
            "name": self.config.backend_name,
            "type": "webdav",
            "stability_class": "stable",
            "read_priority": 80,
            "write_priority": 80,
            "config": {
                "endpoint_url": self.config.openlist_dav_url,
                "username": self.config.openlist_username,
                "password": self.config.openlist_password,
                "root_path": self.config.openlist_root_path,
                "verify_ssl": False,
            },
        }
        response = self._request("POST", "/api/backends", json_body=payload)
        return response["backend"]

    def list_backends(self) -> list[dict[str, object]]:
        return self._request("GET", "/api/backends")["items"]

    def check_backend(self, backend_id: str) -> dict[str, object]:
        return self._request("POST", f"/api/backends/{backend_id}/check")

    def disable_backend(self, backend_id: str) -> dict[str, object]:
        return self._request("POST", f"/api/backends/{backend_id}/disable")["backend"]

    def enable_backend(self, backend_id: str) -> dict[str, object]:
        return self._request("POST", f"/api/backends/{backend_id}/enable")["backend"]

    def list_policies(self) -> list[dict[str, object]]:
        return self._request("GET", "/api/policies/placement")["items"]

    def list_jobs(self) -> list[dict[str, object]]:
        return self._request("GET", "/api/jobs")["items"]

    def upsert_policy(
        self,
        file_class: str,
        *,
        require_local: bool,
        stable_replica_count: int,
        opportunistic_replica_count: int,
        preferred_backend_ids: list[str],
        excluded_backend_ids: list[str],
        max_non_local_size_bytes: int | None,
    ) -> dict[str, object]:
        payload = {
            "require_local": require_local,
            "stable_replica_count": stable_replica_count,
            "opportunistic_replica_count": opportunistic_replica_count,
            "preferred_backend_ids": preferred_backend_ids,
            "excluded_backend_ids": excluded_backend_ids,
            "max_non_local_size_bytes": max_non_local_size_bytes,
        }
        return self._request("PUT", f"/api/policies/placement/{file_class}", json_body=payload)["policy"]

    def upload_text_file(self) -> str:
        return self.upload_bytes(
            filename=self.config.test_filename,
            content=self.config.test_content.encode("utf-8"),
            directory_id=self.config.directory_id,
        )

    def upload_bytes(self, *, filename: str, content: bytes, directory_id: str) -> str:
        metadata_value = base64.b64encode(filename.encode("utf-8")).decode("ascii")
        headers = {
            "Tus-Resumable": "1.0.0",
            "Upload-Length": str(len(content)),
            "Upload-Metadata": f"filename {metadata_value}",
        }
        create_response, create_headers = self._raw_request("POST", "/files", headers=headers)
        _ = create_response
        location = header_value(create_headers, "Location")
        upload_id = location.rsplit("/", 1)[-1]
        patch_headers = {
            "Tus-Resumable": "1.0.0",
            "Upload-Offset": "0",
            "Content-Type": "application/offset+octet-stream",
        }
        self._raw_request("PATCH", f"/files/{upload_id}", headers=patch_headers, data=content)
        finalize = self._request(
            "POST",
            f"/api/uploads/{upload_id}/finalize",
            json_body={"directory_id": directory_id, "filename": filename},
        )
        return finalize["file"]["id"]

    def wait_for_ready_replicas(self, file_id: str, backend_id: str) -> dict[str, object]:
        return self.wait_for_ready_replica_set(file_id, {"bkd_local_default", backend_id})

    def wait_for_ready_replica_set(self, file_id: str, backend_ids: set[str]) -> dict[str, object]:
        deadline = time.time() + self.config.reconcile_timeout_seconds
        detail: dict[str, object] = {}
        while time.time() < deadline:
            self.reconcile_file(file_id)
            self.run_pending_jobs()
            detail = self.file_detail(file_id)
            ready_ids = set(detail["ready_replica_backend_ids"])
            if backend_ids.issubset(ready_ids):
                return detail
            time.sleep(self.config.reconcile_poll_interval_seconds)
        raise RuntimeError(f"timed out waiting for replicas: {json.dumps(detail, ensure_ascii=False)}")

    def download_file(self, file_id: str) -> bytes:
        body, _headers = self._raw_request("GET", f"/api/files/{file_id}/download")
        return body

    def file_detail(self, file_id: str) -> dict[str, object]:
        return self._request("GET", f"/api/files/{file_id}")

    def stream_file(self, file_id: str, *, range_header: str | None = None) -> tuple[bytes, dict[str, str]]:
        headers = {"Range": range_header} if range_header else {}
        return self._raw_request("GET", f"/api/files/{file_id}/stream", headers=headers)

    def delete_file(self, file_id: str) -> dict[str, object]:
        return self._request("DELETE", f"/api/files/{file_id}")

    def restore_file(self, file_id: str) -> dict[str, object]:
        return self._request("POST", f"/api/files/{file_id}/restore")

    def list_recycle_bin(self) -> list[dict[str, object]]:
        return self._request("GET", "/api/files/recycle-bin")["items"]

    def reconcile_file(self, file_id: str) -> dict[str, object]:
        return self._request("POST", f"/api/files/{file_id}/reconcile")

    def run_pending_jobs(self) -> dict[str, object]:
        return self._request("POST", "/api/jobs/run-pending")

    def _request(self, method: str, path: str, json_body: dict[str, object] | None = None) -> dict[str, object]:
        payload = None if json_body is None else json.dumps(json_body).encode("utf-8")
        headers = {"Content-Type": "application/json"} if payload is not None else {}
        body, _response_headers = self._raw_request(method, path, headers=headers, data=payload)
        return json.loads(body.decode("utf-8")) if body else {}

    def _raw_request(
        self,
        method: str,
        path: str,
        headers: dict[str, str] | None = None,
        data: bytes | None = None,
    ) -> tuple[bytes, dict[str, str]]:
        merged_headers = dict(headers or {})
        if self.token:
            merged_headers["Authorization"] = f"Bearer {self.token}"
        request = urllib.request.Request(
            self.config.iron_base_url + path,
            method=method,
            data=data,
            headers=merged_headers,
        )
        try:
            with urllib.request.urlopen(request, timeout=self.config.http_timeout_seconds) as response:
                return response.read(), dict(response.headers.items())
        except urllib.error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"{method} {path} failed with {exc.code}: {body}") from exc


def read_openlist_webdav(endpoint_url: str, username: str, password: str, relative_path: str) -> bytes:
    token = base64.b64encode(f"{username}:{password}".encode("utf-8")).decode("ascii")
    path = "/".join(urllib.parse.quote(part, safe="") for part in relative_path.strip("/").split("/") if part)
    request = urllib.request.Request(
        f"{endpoint_url.rstrip('/')}/{path}",
        method="GET",
        headers={"Authorization": f"Basic {token}"},
    )
    with urllib.request.urlopen(request, timeout=30) as response:
        return response.read()


def header_value(headers: dict[str, str], name: str) -> str:
    for key, value in headers.items():
        if key.lower() == name.lower():
            return value
    raise KeyError(name)


if __name__ == "__main__":
    raise SystemExit(main())