"""pull の回収ホワイトリスト（SPEC_01「回収ホワイトリスト」）の判定と，受け取った tar の検査.
回収できるのは `runs/<RUN_ID>/` 配下で名前が一致するファイルのみ．拡張子や中身では判断せず，
列挙外はすべて回収しない．判定は 2 か所で使う（二重の防御．SPEC_03「pull」）．
1. サーバの一覧（`build_pull_list_command` の出力）から送らせるファイルを絞り込む（`plan_pull`）
2. 手元で受け取った tar を展開前に全メンバー検査する（`check_pull_members`）
エラーメッセージ・戻り値の表示用情報には件数だけを載せ，ファイル名は出さない
（患者情報を含む名前の可能性があるため）．
"""
from __future__ import annotations
import re
import tarfile
from collections.abc import Iterable
from dataclasses import dataclass
# runs/<RUN_ID>/ 直下で回収するファイル名
PULL_ROOT_FILES = frozenset(
{
"config.yaml",
"history.csv",
"metrics.json",
"confusion_matrix.csv",
"log.txt",
"dataset_summary.json",
}
)
PULL_FIGURES_DIR = "figures"
PULL_FIGURE_SUFFIX = ".png"
PULL_CHECKPOINTS_DIR = "checkpoints"
PULL_CHECKPOINT_SUFFIX = ".pt"
# find -printf '%y' の通常ファイル
_FIND_TYPE_FILE = "f"
_WINDOWS_DRIVE_RE = re.compile(r"^[A-Za-z]:")
@dataclass(frozen=True)
class PullPlan:
"""サーバから送らせるファイルの計画.
Attributes:
files: 送らせるファイル（`runs/` 基準の POSIX パス．表示しない）．
skipped: 回収しなかったファイルの件数（リンク等の通常ファイル以外を含む）．
skipped_checkpoints: そのうち `--with-checkpoints` を付ければ回収できた件数．
"""
files: tuple[str, ...]
skipped: int
skipped_checkpoints: int
def validate_run_id(run_id: str) -> str:
"""RUN_ID（`runs/` 直下のディレクトリ名）を検証して返す.
Args:
run_id: 実験 ID．
Returns:
検証済みの RUN_ID．
Raises:
ValueError: 空・`/`・`\\`・`..`・制御文字を含む，`.`/`-` で始まる場合．
"""
if (
not run_id
or "/" in run_id
or "\\" in run_id
or ".." in run_id
or run_id.startswith((".", "-"))
or any(ord(c) < 0x20 or c == "\x7f" for c in run_id)
or run_id != run_id.strip()
):
raise ValueError(
f"RUN_ID が不正である: {run_id!r}（runs/ 直下のディレクトリ名を指定する．"
"/・\\・.. は使えない）"
)
return run_id
def _is_safe_part(part: str) -> bool:
"""パスの 1 要素が `..`・`.`・空・バックスラッシュ・ドライブ文字を含まないか."""
return part not in ("", ".", "..") and "\\" not in part and not _WINDOWS_DRIVE_RE.match(part)
def _is_safe_relative(path: str) -> bool:
"""POSIX 区切りの相対パスとして安全か（絶対パス・パス外へ出る要素を含まない）."""
return not path.startswith("/") and all(_is_safe_part(part) for part in path.split("/"))
def is_pull_allowed(rel_path: str, *, with_checkpoints: bool) -> bool:
"""`runs/<RUN_ID>/` 基準の相対パスが回収ホワイトリストに一致するかを返す.
Args:
rel_path: `runs/<RUN_ID>/` 基準の POSIX 相対パス．
with_checkpoints: `checkpoints/*.pt` も回収対象にするか．
Returns:
回収してよければ True．
"""
if not _is_safe_relative(rel_path):
return False
parts = rel_path.split("/")
if len(parts) == 1:
return parts[0] in PULL_ROOT_FILES
if len(parts) != 2:
return False
directory, name = parts
if directory == PULL_FIGURES_DIR:
return name.endswith(PULL_FIGURE_SUFFIX) and len(name) > len(PULL_FIGURE_SUFFIX)
if directory == PULL_CHECKPOINTS_DIR and with_checkpoints:
return name.endswith(PULL_CHECKPOINT_SUFFIX) and len(name) > len(PULL_CHECKPOINT_SUFFIX)
return False
def _strip_run_prefix(path: str, run_id: str) -> str | None:
"""`<RUN_ID>/<rel>` から `<rel>` を取り出す（RUN_ID 配下でなければ None）."""
prefix = f"{run_id}/"
return path[len(prefix) :] if path.startswith(prefix) else None
def parse_find_listing(data: bytes) -> list[tuple[str, str]]:
"""`build_pull_list_command` の出力（`<種類> <パス>` の NUL 区切り）を分解する.
Linux のファイル名は任意のバイト列なので，往復で失われないよう surrogateescape で扱う．
Args:
data: サーバの出力．
Returns:
(種類, `runs/` 基準のパス) の一覧．
Raises:
ValueError: 書式が崩れている場合（件数のみ示す）．
"""
entries: list[tuple[str, str]] = []
for number, record in enumerate((r for r in data.split(b"\0") if r), start=1):
kind, sep, path = record.decode("utf-8", "surrogateescape").partition(" ")
if not sep or not kind or not path:
raise ValueError(f"サーバのファイル一覧の {number} 件目を解釈できない")
entries.append((kind, path))
return entries
def plan_pull(
listing: Iterable[tuple[str, str]], run_id: str, *, with_checkpoints: bool
) -> PullPlan:
"""サーバのファイル一覧から，送らせるファイル（ホワイトリストに一致する通常ファイル）を選ぶ.
Args:
listing: `parse_find_listing` の結果．
run_id: 検証済みの RUN_ID．
with_checkpoints: `checkpoints/*.pt` も回収するか．
Returns:
送らせるファイルと，回収しなかった件数．
"""
files: list[str] = []
skipped = 0
skipped_checkpoints = 0
for kind, path in listing:
rel = _strip_run_prefix(path, run_id)
if kind == _FIND_TYPE_FILE and rel is not None:
if is_pull_allowed(rel, with_checkpoints=with_checkpoints):
files.append(path)
continue
if not with_checkpoints and is_pull_allowed(rel, with_checkpoints=True):
skipped_checkpoints += 1
skipped += 1
return PullPlan(tuple(sorted(files)), skipped, skipped_checkpoints)
def check_pull_members(
members: Iterable[tarfile.TarInfo], run_id: str, *, with_checkpoints: bool
) -> int:
"""受け取った tar の全メンバーを展開前に検査する.
1 つでも問題があれば例外を送出する（呼び出し側は何も展開しない）．受け入れるのは
`<RUN_ID>/<ホワイトリストに一致する名前>` の通常ファイルのみで，ディレクトリ・リンク
（シンボリック・ハード）・特殊ファイルは拒否する（サーバはファイル名だけを tar に渡すため，
正常な応答にディレクトリのメンバーは含まれない）．
Args:
members: tar のメンバー（`TarFile.getmembers()`）．
run_id: 検証済みの RUN_ID．
with_checkpoints: `checkpoints/*.pt` を受け入れるか．
Returns:
受け入れたメンバーの件数．
Raises:
ValueError: ホワイトリスト外の名前・パス外に出る名前（`..`・絶対パス・ドライブ文字・
バックスラッシュ）・リンク・通常ファイル以外が含まれる場合（件数のみを示す）．
"""
accepted = 0
outside = 0
links = 0
not_regular = 0
not_allowed = 0
for member in members:
if not _is_safe_relative(member.name):
outside += 1
elif member.issym() or member.islnk():
links += 1
elif not member.isreg():
not_regular += 1
else:
rel = _strip_run_prefix(member.name, run_id)
if rel is None or not is_pull_allowed(rel, with_checkpoints=with_checkpoints):
not_allowed += 1
else:
accepted += 1
if outside or links or not_regular or not_allowed:
raise ValueError(
"受け取ったアーカイブに回収できないメンバーが含まれるため，何も展開せず中止した"
f"（ホワイトリスト外 {not_allowed} 件・パス外 {outside} 件・リンク {links} 件・"
f"通常ファイル以外 {not_regular} 件）"
)
return accepted
