Newer
Older
TIASDataHub / src / tiasdata / store.py
"""ストア(登録したファイルの実体)のハッシュ・コピー・照合.

ストアのファイルは中身の SHA-256 を名前にし,一度置いたら書き換えない・消さない
(SPEC_01 基本方針 2「登録したファイルは変えない」).置き場所は
`store/<SHA-256 の先頭 2 文字>/<SHA-256>.<拡張子>` で中身から決まるため,目録には保存しない
(SPEC_02「assets」).

例外メッセージ・ログにはパス(データの場所)や元のファイル名を出さず,SHA-256 で示す.
"""

import hashlib
import logging
import os
import re
import stat
import tempfile
from enum import StrEnum
from pathlib import Path
from typing import NamedTuple

logger = logging.getLogger(__name__)

# 大きいファイルでもメモリを使い切らないよう分割して読む
CHUNK_SIZE = 1024 * 1024

# 書きかけの一時ファイルの名前の接頭辞.正式な名前(<sha256>.<ext>)と区別でき,
# 点検(`tiasdata check`)で書きかけの残りを見分けられるようにする
TEMP_PREFIX = ".tmp-"

_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$")
_EXT_PATTERN = re.compile(r"^[a-z0-9]+$")


class StoredFile(NamedTuple):
    """ストアへの取り込みの結果.

    Attributes:
        sha256: 中身の SHA-256(小文字の 16 進 64 文字).
        ext: 拡張子(小文字.先頭の `.` なし).
        size_bytes: サイズ.
        copied: 今回コピーしたか(同じ中身が既にあったときは False).
    """

    sha256: str
    ext: str
    size_bytes: int
    copied: bool


class StoredFileStatus(StrEnum):
    """ストアのファイルの照合の結果."""

    OK = "ok"
    MISSING = "missing"
    MISMATCH = "mismatch"


def normalize_ext(ext: str) -> str:
    """拡張子を小文字・先頭の `.` なしの形にそろえる.

    Args:
        ext: 拡張子(`.PNG`・`png` 等).

    Returns:
        小文字の拡張子(例: `png`).

    Raises:
        ValueError: 空,または英数字以外を含む.
    """
    normalized = ext.removeprefix(".").lower()
    if not _EXT_PATTERN.match(normalized):
        # 拡張子そのものは出さない(元のファイル名の一部のため)
        raise ValueError("拡張子が無い,または英数字以外を含むファイルはストアに置けない")
    return normalized


def store_path(store_dir: Path, sha256: str, ext: str) -> Path:
    """ストア上のパスを決める(ファイルには触れない純関数).

    Args:
        store_dir: ストアのフォルダ(`<catalog.root>/store`).
        sha256: 中身の SHA-256(小文字の 16 進 64 文字).
        ext: 拡張子(`.PNG` 等も小文字・先頭の `.` なしにそろえる).

    Returns:
        `store_dir/<sha256 の先頭 2 文字>/<sha256>.<ext>`.

    Raises:
        ValueError: SHA-256 の形が誤っている,または拡張子が不正.
    """
    if not _SHA256_PATTERN.match(sha256):
        raise ValueError("SHA-256 は小文字の 16 進 64 文字で指定する")
    return store_dir / sha256[:2] / f"{sha256}.{normalize_ext(ext)}"


def sha256_file(path: Path) -> str:
    """ファイルの SHA-256 を計算する(読み取り専用で開き,分割して読む).

    Args:
        path: 対象のファイル.

    Returns:
        SHA-256(小文字の 16 進 64 文字).
    """
    digest = hashlib.sha256()
    with path.open("rb") as f:
        while chunk := f.read(CHUNK_SIZE):
            digest.update(chunk)
    return digest.hexdigest()


def verify_stored(store_dir: Path, sha256: str, ext: str) -> StoredFileStatus:
    """ストアのファイルが期待する SHA-256 と一致するかを照合する(点検用).

    Args:
        store_dir: ストアのフォルダ.
        sha256: 期待する SHA-256.
        ext: 拡張子.

    Returns:
        一致すれば `OK`,ファイルが無ければ `MISSING`,中身が違えば `MISMATCH`.

    Raises:
        ValueError: SHA-256 の形が誤っている,または拡張子が不正.
    """
    path = store_path(store_dir, sha256, ext)
    if not path.is_file():
        return StoredFileStatus.MISSING
    if sha256_file(path) != sha256:
        return StoredFileStatus.MISMATCH
    return StoredFileStatus.OK


def add_file(store_dir: Path, source: Path) -> StoredFile:
    """取り込み元のファイルをストアに置く.

    取り込み元は読み取り専用で開くだけで,中身・更新日時・属性を変えない.
    同じ中身が既にあれば(同じパスにあり照合が合えば)コピーしない.
    無ければ同じフォルダの一時ファイルへコピーし,書き込んだファイルの SHA-256 を照合してから
    正式な名前に変える(書きかけのファイルが正式な名前で残らないようにするため).
    置いたファイルには読み取り専用の属性を付ける.既存のファイルは上書きしない
    (照合が合って再利用するときは,属性だけを読み取り専用に付け直す).

    Args:
        store_dir: ストアのフォルダ(`tiasdata init` で作成済みであること).
        source: 取り込み元のファイル.

    Returns:
        SHA-256・拡張子・サイズ・今回コピーしたか.

    Raises:
        ValueError: ストアのフォルダが無い,取り込み元がファイルでない,拡張子が不正,
            コピー中に取り込み元が変わった,コピーしたファイルの照合が合わない,
            または同じパスに中身の違うファイルが既にある.
    """
    if not store_dir.is_dir():
        raise ValueError("ストアのフォルダが無い(先に tiasdata init を実行する)")
    if not source.is_file():
        raise ValueError("取り込み元のファイルが無い,またはファイルではない")
    ext = normalize_ext(source.suffix)

    sha256 = sha256_file(source)
    target = store_path(store_dir, sha256, ext)
    if target.exists():
        return _reuse_existing(target, sha256, ext)

    target.parent.mkdir(exist_ok=True)
    tmp_path, copied_sha256, size_bytes = _copy_to_temp(source, target.parent)
    try:
        # 1 回目の計算からコピーまでの間に取り込み元が変わっていないか
        if copied_sha256 != sha256:
            raise ValueError(f"コピー中に取り込み元の中身が変わった(SHA-256 {sha256})")
        # 書き込んだファイルを読み直して照合する(ディスクへの書き込みの誤りを検出するため)
        if sha256_file(tmp_path) != sha256:
            raise ValueError(f"ストアにコピーしたファイルの照合が合わない(SHA-256 {sha256})")
        published = _publish(tmp_path, target)
    finally:
        tmp_path.unlink(missing_ok=True)

    if not published:
        # 照合の間に同じ名前のファイルが置かれた(別の取り込みと同時に実行した等)
        return _reuse_existing(target, sha256, ext)
    _make_read_only(target)
    logger.debug("ストアに追加した: %s", sha256)
    return StoredFile(sha256=sha256, ext=ext, size_bytes=size_bytes, copied=True)


def _reuse_existing(target: Path, sha256: str, ext: str) -> StoredFile:
    # 同じ名前のファイルは同じ中身のはず.違えばストアが壊れているので上書きせず止める
    if not target.is_file() or sha256_file(target) != sha256:
        raise ValueError(
            f"ストアに同じ名前で中身の違うファイルがある(SHA-256 {sha256})."
            "上書きせずに停止した.tiasdata check で点検する"
        )
    # 復元・手作業等で読み取り専用が外れた正しいファイルを,取り込みのたびに守り直すため
    _make_read_only(target)
    return StoredFile(sha256=sha256, ext=ext, size_bytes=target.stat().st_size, copied=False)


def _copy_to_temp(source: Path, directory: Path) -> tuple[Path, str, int]:
    # 一時ファイルは正式な名前と同じフォルダに作る(同じドライブ内の rename にするため)
    digest = hashlib.sha256()
    size_bytes = 0
    fd, tmp_name = tempfile.mkstemp(prefix=TEMP_PREFIX, dir=directory)
    tmp_path = Path(tmp_name)
    try:
        with os.fdopen(fd, "wb") as dst, source.open("rb") as src:
            while chunk := src.read(CHUNK_SIZE):
                digest.update(chunk)
                dst.write(chunk)
                size_bytes += len(chunk)
            dst.flush()
            os.fsync(dst.fileno())
    except BaseException:
        tmp_path.unlink(missing_ok=True)
        raise
    return tmp_path, digest.hexdigest(), size_bytes


def _publish(tmp_path: Path, target: Path) -> bool:
    """一時ファイルを正式な名前にする.既に同じ名前があれば何もせず False を返す."""
    try:
        if os.name == "nt":
            # Windows の rename は移動先があると FileExistsError になり,上書きしない
            os.rename(tmp_path, target)
        else:
            # POSIX の rename は黙って上書きするため,既存があれば失敗する link を使う
            os.link(tmp_path, target)
    except FileExistsError:
        return False
    return True


def _make_read_only(path: Path) -> None:
    # 書き換え・削除の誤操作を防ぐ(Windows では読み取り専用の属性になる)
    mode = path.stat().st_mode
    path.chmod(mode & ~(stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH))


class StoreEntry(NamedTuple):
    """ストアの正式な名前のファイル 1 つ(点検用).

    Attributes:
        sha256: ファイル名の SHA-256(中身と一致するかは確かめていない).
        ext: 拡張子.
        writable: 読み取り専用の属性が外れているか.
    """

    sha256: str
    ext: str
    writable: bool


class StoreScan(NamedTuple):
    """ストアのフォルダを走査した結果(点検用).

    Attributes:
        entries: 正式な名前(`<先頭 2 文字>/<sha256>.<ext>`)のファイル.
        temp_files: 書きかけの一時ファイル(`TEMP_PREFIX` で始まる名前)の数.
        unexpected: 上のどちらにも当たらないファイル・フォルダの数.
    """

    entries: list[StoreEntry]
    temp_files: int
    unexpected: int


_STORED_NAME_PATTERN = re.compile(r"^([0-9a-f]{64})\.([a-z0-9]+)$")
_PREFIX_DIR_PATTERN = re.compile(r"^[0-9a-f]{2}$")


def scan_store(store_dir: Path) -> StoreScan:
    """ストアのフォルダを走査し,ファイルを正式な名前・一時ファイル・想定外に分ける.

    中身は読まない(ハッシュの照合は `verify_stored` で行う).

    Args:
        store_dir: ストアのフォルダ.

    Returns:
        走査の結果(正式な名前のファイルは SHA-256 の順).

    Raises:
        ValueError: ストアのフォルダが無い.
    """
    if not store_dir.is_dir():
        raise ValueError("ストアのフォルダが無い(先に tiasdata init を実行する)")
    entries: list[StoreEntry] = []
    temp_files = 0
    unexpected = 0
    for prefix_dir in store_dir.iterdir():
        if not prefix_dir.is_dir() or not _PREFIX_DIR_PATTERN.match(prefix_dir.name):
            unexpected += 1
            continue
        for path in prefix_dir.iterdir():
            if path.name.startswith(TEMP_PREFIX):
                temp_files += 1
                continue
            match = _STORED_NAME_PATTERN.match(path.name)
            if not path.is_file() or match is None or match.group(1)[:2] != prefix_dir.name:
                unexpected += 1
                continue
            entries.append(StoreEntry(match.group(1), match.group(2), is_writable(path)))
    entries.sort()
    return StoreScan(entries=entries, temp_files=temp_files, unexpected=unexpected)


def is_writable(path: Path) -> bool:
    """ストアのファイルの読み取り専用の属性が外れているか(所有者の書き込み権限があるか).

    Args:
        path: ストアのファイル.

    Returns:
        書き込みできる属性なら True(Windows では読み取り専用の属性が無いとき).
    """
    return bool(path.stat().st_mode & stat.S_IWUSR)