Newer
Older
SmTIAS-LightSim / src / sm_tias_lightsim / remote / transfer.py
"""tar ストリームの書き出しと展開(ssh の標準入出力につなぐ).

Windows には rsync が無いため,Python の tarfile でストリームを作り
`ssh <host> 'tar -x ...'` の標準入力へ流す(push),または `ssh <host> 'tar -c ...'` の
標準出力を展開する(pull).手元の tar に依存しないので PowerShell からも Git Bash からも
同じに動く.

- 書き出しは PAX 形式(日本語ファイル名を UTF-8 で確実に運ぶ)
- 展開は,メンバー名の構造検証(`..`・`.`・空の要素・バックスラッシュ・ドライブ文字を拒否)→
  リンク(シンボリック・ハード)の一律拒否 → 要求したパスとの前方一致 → Python 3.12 の
  `filter="data"`(展開先の外・特殊ファイル)の順に確かめ,どれかに当たれば中断する
"""

from __future__ import annotations

import re
import tarfile
from collections.abc import Callable, Sequence
from pathlib import Path
from typing import BinaryIO

from sm_tias_lightsim.remote.files import PushEntry

# (処理済みバイト数, 総バイト数.不明なら 0)
ProgressCallback = Callable[[int, int], None]

_ENCODING = "utf-8"
# Windows のドライブ指定(`C:` 形式).os.path.join がこれを見ると展開先を無視して付け替える
_WINDOWS_DRIVE_RE = re.compile(r"^[A-Za-z]:")


def _normalize_tarinfo(info: tarfile.TarInfo) -> tarfile.TarInfo:
    """所有者とパーミッションを送信元 OS に依存しない値にそろえる.

    Windows の stat は所有者を持たず実行ビットも拡張子から推測した値なので,
    展開結果がそれに引きずられないよう固定する(実行ビットが立っているものだけ 755).
    """
    info.uid = info.gid = 0
    info.uname = info.gname = ""
    info.mode = 0o755 if info.isdir() or info.mode & 0o111 else 0o644
    return info


def write_tar_stream(
    entries: Sequence[PushEntry],
    fileobj: BinaryIO,
    *,
    compress: bool,
    progress: ProgressCallback | None = None,
) -> int:
    """送信対象を tar ストリームとして fileobj へ書き出す.

    fileobj は閉じない(呼び出し側が ssh の標準入力を閉じて EOF を伝える).

    Args:
        entries: 送信対象.
        fileobj: 書き込み先のバイナリストリーム(ssh の標準入力).
        compress: True なら gzip 圧縮する.
        progress: ファイル 1 件ごとに呼ぶ進捗コールバック(任意).

    Returns:
        書き出したファイルの合計バイト数(圧縮前).

    Raises:
        OSError: 書き込み先が閉じられた(ssh が異常終了した)場合など.
    """
    total = sum(entry.size for entry in entries)
    done = 0
    mode = "w|gz" if compress else "w|"
    with tarfile.open(
        fileobj=fileobj,
        mode=mode,
        format=tarfile.PAX_FORMAT,
        encoding=_ENCODING,
        dereference=True,
    ) as tar:
        for entry in entries:
            tar.add(entry.path, arcname=entry.arcname, recursive=False, filter=_normalize_tarinfo)
            done += entry.size
            if progress is not None:
                progress(done, total)
    return done


def _reject(reason: str, name: str) -> tarfile.FilterError:
    """展開を中止するための例外を作る(問題のメンバー名と確認先を必ず含める)."""
    return tarfile.FilterError(
        f"{reason}: 「{name}」(展開を中止しました.サーバ側の output/ を確認してください)"
    )


def _validate_member_name(name: str) -> None:
    """メンバー名の構造を検証する(要求パスとの前方一致判定の前提).

    `filter="data"` は展開先ディレクトリの外に出るものしか拒否しないため,
    `output/foo/../../src/...` のように展開先の内側で要求パスの外へ出る名前が通ってしまう.
    文字列の前方一致だけでは防げないので,要素単位で危険な形を拒否する.
    バックスラッシュは Windows の `os.path.join` が区切りとして解釈するため
    (`output/foo/..\\..\\src\\x.py` で抜けられる),`posixpath.normpath` で丸める方式は採らない.

    Raises:
        tarfile.FilterError: `..`・`.`・空の要素(先頭 `/` や `//`)・バックスラッシュ・
            ドライブ文字(`C:` 形式)を含む場合.
    """
    # tarfile はディレクトリ名の末尾の / を落とすが,念のため 1 つだけ許す
    trimmed = name[:-1] if name.endswith("/") else name
    for part in trimmed.split("/"):
        if part in ("", ".", "..") or "\\" in part or _WINDOWS_DRIVE_RE.match(part):
            raise _reject(
                "アーカイブのメンバー名が不正です"
                "('..'・'.'・空の要素・バックスラッシュ・ドライブ文字は受け付けません)",
                name,
            )


def _is_under(name: str, prefixes: Sequence[str]) -> bool:
    """name が prefixes のいずれか自身,またはその配下か(構造検証済みの名前に使う)."""
    name = name.rstrip("/")
    return any(name == prefix or name.startswith(prefix + "/") for prefix in prefixes)


def extract_tar_stream(
    fileobj: BinaryIO,
    dest: Path,
    allowed_prefixes: Sequence[str],
    *,
    progress: ProgressCallback | None = None,
) -> tuple[int, int]:
    """tar ストリーム(非圧縮・gzip 自動判別)を dest へ安全に展開する.

    アーカイブの末尾(終端ブロック後のパディング)は読み残すことがあるため,
    ssh プロセスを待つ前に呼び出し側で fileobj を読み切ること(読み残すと ssh が
    パイプへの書き込みで止まり,待ち合わせがデッドロックする).

    Args:
        fileobj: 読み出し元のバイナリストリーム(ssh の標準出力).
        dest: 展開先ディレクトリ(無ければ作る).
        allowed_prefixes: 受け入れるメンバーのパス(要求したパスそのもの).
        progress: ファイル 1 件ごとに呼ぶ進捗コールバック(任意.総バイト数は 0).

    Returns:
        (展開したファイル数, 合計バイト数).

    検証の順序: メンバー名の構造 → リンクの拒否 → 要求パスとの前方一致 → `data_filter`
    (展開先の外・特殊ファイル・危険なパーミッション).中断時,それまでに展開した
    ファイルは残る.

    Raises:
        tarfile.TarError: ストリームが空・壊れている,または不正な名前・リンク・要求外・
            危険なメンバーを含む場合(名前・リンク・要求外は `tarfile.FilterError`).
        OSError: 書き込みに失敗した場合.
    """
    counts = {"files": 0, "bytes": 0}

    def _filter(member: tarfile.TarInfo, path: str) -> tarfile.TarInfo | None:
        _validate_member_name(member.name)
        # data_filter はリンク先が展開先の内側なら通すため,同名の通常ファイルを後から
        # 送られるとハードリンク経由で src/ 等を上書きできる(Windows でも特権なしで作れる).
        # pull の対象(output/ の PNG・npy・csv)にリンクは無い前提で一律に拒否する
        if member.issym() or member.islnk():
            raise _reject(
                f"アーカイブにリンクが含まれています(リンク先 {member.linkname})", member.name
            )
        if not _is_under(member.name, allowed_prefixes):
            raise _reject("要求していないパスがアーカイブに含まれています", member.name)
        filtered = tarfile.data_filter(member, path)
        if filtered is not None and filtered.isfile():
            counts["files"] += 1
            counts["bytes"] += filtered.size
            if progress is not None:
                progress(counts["bytes"], 0)
        return filtered

    dest.mkdir(parents=True, exist_ok=True)
    with tarfile.open(fileobj=fileobj, mode="r|*", encoding=_ENCODING) as tar:
        tar.extractall(dest, filter=_filter)
    return counts["files"], counts["bytes"]