UnifyDrive2NAS/unifydrive2nas.py

#!/usr/bin/env python3
"""Unify Drive の画像・動画を日付別に NAS へコピーする。"""

import argparse
import contextlib
import datetime
import fcntl
import hashlib
import logging
import os
from pathlib import Path
import shutil
import sqlite3
import struct
import sys
import tempfile
from dataclasses import dataclass
from typing import BinaryIO, Callable, Iterator, List, Optional, Sequence, Tuple

# exifread は内部で絶対 import を使うため、同梱先を import 検索パスに加える。
VENDOR_DIRECTORY = Path(__file__).resolve().parent / "vendor"
sys.path.insert(0, str(VENDOR_DIRECTORY))
import exifread  # noqa: E402

logging.getLogger("exifread").setLevel(logging.ERROR)


DEFAULT_SOURCE = Path("/volume1/homes/kazu634/UnifyDrive")
DEFAULT_DESTINATION = Path("/volume1/homes/kazu634/Photos/PhotoLibrary")
DEFAULT_DATABASE = Path.home() / ".local" / "state" / "unifydrive2nas" / "copied.sqlite3"
SUPPORTED_SUFFIXES = frozenset((".jpg", ".jpeg", ".heic", ".heif", ".png", ".mov", ".mp4"))
VIDEO_SUFFIXES = frozenset((".mov", ".mp4"))
JST = datetime.timezone(datetime.timedelta(hours=9))
UTC = datetime.timezone.utc
QUICKTIME_EPOCH = datetime.datetime(1904, 1, 1, tzinfo=UTC)
APPLE_CREATION_KEY = b"com.apple.quicktime.creationdate"

LOGGER = logging.getLogger("unifydrive2nas")


@dataclass
class Config:
    source: Path
    destination: Path
    database: Path
    initialize: bool = False
    dry_run: bool = False


@dataclass
class Summary:
    targets: int = 0
    new: int = 0
    copied: int = 0
    identical: int = 0
    numbered: int = 0
    failed: int = 0


@dataclass(frozen=True)
class Atom:
    kind: bytes
    payload_start: int
    end: int


def iter_target_files(source: Path, on_error: Optional[Callable[[OSError], None]] = None) -> Iterator[Path]:
    """隠し項目と @eaDir を除外し、対象拡張子のファイルを列挙する。"""
    for root, directories, files in os.walk(str(source), topdown=True, onerror=on_error):
        directories[:] = sorted(
            name for name in directories if not name.startswith(".") and name != "@eaDir"
        )
        for name in sorted(files):
            if name.startswith("."):
                continue
            path = Path(root) / name
            if path.suffix.lower() in SUPPORTED_SUFFIXES:
                yield path


def parse_exif_datetime(path: Path) -> Optional[datetime.datetime]:
    """EXIF の原日時、デジタル化日時の順に JST として読む。"""
    with path.open("rb") as handle:
        tags = exifread.process_file(handle, details=False)
    for key in ("EXIF DateTimeOriginal", "EXIF DateTimeDigitized"):
        tag = tags.get(key)
        if tag is None:
            continue
        try:
            value = datetime.datetime.strptime(str(tag).strip(), "%Y:%m:%d %H:%M:%S")
            return value.replace(tzinfo=JST)
        except (TypeError, ValueError):
            continue
    return None


def _iter_atoms(handle: BinaryIO, start: int, end: int) -> Iterator[Atom]:
    """指定範囲の ISO Base Media File Format atom を安全に走査する。"""
    position = start
    while position + 8 <= end:
        handle.seek(position)
        header = handle.read(8)
        if len(header) != 8:
            return
        size, kind = struct.unpack(">I4s", header)
        header_size = 8
        if size == 1:
            extended = handle.read(8)
            if len(extended) != 8:
                return
            size = struct.unpack(">Q", extended)[0]
            header_size = 16
        elif size == 0:
            size = end - position
        if size < header_size or position + size > end:
            return
        yield Atom(kind=kind, payload_start=position + header_size, end=position + size)
        position += size


def _find_child(handle: BinaryIO, start: int, end: int, kind: bytes) -> Optional[Atom]:
    for atom in _iter_atoms(handle, start, end):
        if atom.kind == kind:
            return atom
    return None


def _read_apple_creation_date(handle: BinaryIO, meta: Atom) -> Optional[datetime.datetime]:
    # QuickTime 形式は payload 直下、ISO BMFF FullBox 形式は version/flags の後に子がある。
    keys = None
    ilst = None
    for children_start in (meta.payload_start, meta.payload_start + 4):
        keys = _find_child(handle, children_start, meta.end, b"keys")
        ilst = _find_child(handle, children_start, meta.end, b"ilst")
        if keys is not None and ilst is not None:
            break
    if keys is None or ilst is None:
        return None

    handle.seek(keys.payload_start)
    keys_header = handle.read(8)
    if len(keys_header) != 8:
        return None
    key_count = struct.unpack(">I", keys_header[4:8])[0]
    wanted_index = None
    position = keys.payload_start + 8
    for index in range(1, key_count + 1):
        if position + 8 > keys.end:
            return None
        handle.seek(position)
        entry_header = handle.read(8)
        entry_size, namespace = struct.unpack(">I4s", entry_header)
        if entry_size < 8 or position + entry_size > keys.end:
            return None
        key = handle.read(entry_size - 8)
        if namespace == b"mdta" and key == APPLE_CREATION_KEY:
            wanted_index = index
            break
        position += entry_size
    if wanted_index is None:
        return None

    item = _find_child(handle, ilst.payload_start, ilst.end, struct.pack(">I", wanted_index))
    if item is None:
        return None
    data = _find_child(handle, item.payload_start, item.end, b"data")
    if data is None or data.payload_start + 8 > data.end:
        return None
    handle.seek(data.payload_start + 8)
    raw_value = handle.read(data.end - data.payload_start - 8)
    try:
        text = raw_value.rstrip(b"\x00").decode("utf-8").strip()
    except UnicodeDecodeError:
        return None
    try:
        value = datetime.datetime.strptime(text, "%Y-%m-%dT%H:%M:%S%z")
    except ValueError:
        return None
    return value.astimezone(JST)


def _read_mvhd_datetime(handle: BinaryIO, mvhd: Atom) -> Optional[datetime.datetime]:
    handle.seek(mvhd.payload_start)
    version_bytes = handle.read(4)
    if len(version_bytes) != 4:
        return None
    version = version_bytes[0]
    if version == 0:
        raw = handle.read(4)
        if len(raw) != 4:
            return None
        seconds = struct.unpack(">I", raw)[0]
    elif version == 1:
        raw = handle.read(8)
        if len(raw) != 8:
            return None
        seconds = struct.unpack(">Q", raw)[0]
    else:
        return None
    if seconds == 0:
        return None
    try:
        return (QUICKTIME_EPOCH + datetime.timedelta(seconds=seconds)).astimezone(JST)
    except (OverflowError, ValueError):
        return None


def parse_video_datetimes(path: Path) -> Tuple[Optional[datetime.datetime], Optional[datetime.datetime]]:
    """Apple 作成日時と mvhd 作成日時を読み、両方を返す。"""
    with path.open("rb") as handle:
        file_size = path.stat().st_size
        moov = _find_child(handle, 0, file_size, b"moov")
        if moov is None:
            return None, None
        meta = _find_child(handle, moov.payload_start, moov.end, b"meta")
        mvhd = _find_child(handle, moov.payload_start, moov.end, b"mvhd")
        apple = _read_apple_creation_date(handle, meta) if meta is not None else None
        mvhd_value = _read_mvhd_datetime(handle, mvhd) if mvhd is not None else None
        return apple, mvhd_value


def determine_datetime(path: Path, stat_result: os.stat_result) -> Tuple[datetime.datetime, str]:
    """仕様で定めた優先順で日時と出どころを決める。"""
    try:
        value = parse_exif_datetime(path)
        if value is not None:
            return value, "exif"
    except Exception as exc:  # 壊れたファイルでも後続候補へ進む。
        LOGGER.debug("EXIF の解析に失敗: %s: %s", path, exc)

    if path.suffix.lower() in VIDEO_SUFFIXES:
        try:
            apple, mvhd = parse_video_datetimes(path)
            if apple is not None:
                return apple, "apple"
            if mvhd is not None:
                return mvhd, "mvhd"
        except Exception as exc:  # 壊れた atom でも mtime へ進む。
            LOGGER.debug("動画メタデータの解析に失敗: %s: %s", path, exc)

    return datetime.datetime.fromtimestamp(stat_result.st_mtime, JST), "mtime"


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while True:
            chunk = handle.read(1024 * 1024)
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()


def same_content(left: Path, right: Path) -> bool:
    if left.stat().st_size != right.stat().st_size:
        return False
    return sha256_file(left) == sha256_file(right)


def select_destination(source: Path, directory: Path) -> Tuple[Path, bool, bool]:
    """コピー先、同一内容か、連番名かを返す。"""
    candidate = directory / source.name
    number = 0
    while candidate.exists():
        if candidate.is_file() and same_content(source, candidate):
            return candidate, True, number > 0
        number += 1
        candidate = directory / (source.stem + "_" + str(number) + source.suffix)
    return candidate, False, number > 0


def atomic_copy(source: Path, destination: Path) -> None:
    """同じディレクトリの一時ファイルへ copy2 し、rename で確定する。"""
    file_descriptor, temporary_name = tempfile.mkstemp(
        prefix="." + destination.name + ".tmp-", dir=str(destination.parent)
    )
    os.close(file_descriptor)
    temporary = Path(temporary_name)
    try:
        shutil.copy2(str(source), str(temporary))
        # rename 前に内容とメタデータを確定し、本来の名前に不完全なファイルを出さない。
        with temporary.open("rb") as handle:
            os.fsync(handle.fileno())
        os.replace(str(temporary), str(destination))
    finally:
        try:
            temporary.unlink()
        except FileNotFoundError:
            pass


def open_database(path: Path, read_only: bool = False) -> sqlite3.Connection:
    if read_only and path.exists():
        connection = sqlite3.connect(path.resolve().as_uri() + "?mode=ro", uri=True)
        return connection
    if read_only:
        connection = sqlite3.connect(":memory:")
    else:
        path.parent.mkdir(parents=True, exist_ok=True)
        connection = sqlite3.connect(str(path))
    connection.execute(
        """
        CREATE TABLE IF NOT EXISTS copied_files (
            source_relative_path TEXT NOT NULL,
            source_size INTEGER NOT NULL,
            source_mtime_ns INTEGER NOT NULL,
            destination_path TEXT,
            date_source TEXT NOT NULL,
            action TEXT NOT NULL,
            processed_at TEXT NOT NULL,
            PRIMARY KEY (source_relative_path, source_size, source_mtime_ns)
        )
        """
    )
    connection.commit()
    return connection


def is_recorded(connection: sqlite3.Connection, relative_path: str, size: int, mtime_ns: int) -> bool:
    row = connection.execute(
        """
        SELECT 1 FROM copied_files
        WHERE source_relative_path = ? AND source_size = ? AND source_mtime_ns = ?
        """,
        (relative_path, size, mtime_ns),
    ).fetchone()
    return row is not None


def record_file(
    connection: sqlite3.Connection,
    relative_path: str,
    size: int,
    mtime_ns: int,
    destination_path: Optional[str],
    date_source: str,
    action: str,
) -> None:
    connection.execute(
        """
        INSERT OR REPLACE INTO copied_files (
            source_relative_path, source_size, source_mtime_ns,
            destination_path, date_source, action, processed_at
        ) VALUES (?, ?, ?, ?, ?, ?, ?)
        """,
        (
            relative_path,
            size,
            mtime_ns,
            destination_path,
            date_source,
            action,
            datetime.datetime.now(UTC).isoformat(),
        ),
    )
    connection.commit()


@contextlib.contextmanager
def process_lock(lock_path: Path) -> Iterator[None]:
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    with lock_path.open("a+") as handle:
        try:
            fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            raise RuntimeError("別の処理が実行中です: " + str(lock_path))
        try:
            yield
        finally:
            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)


def validate_source(source: Path, mount_checker: Callable[[str], bool]) -> None:
    if not source.is_dir():
        raise RuntimeError("コピー元がディレクトリではありません: " + str(source))
    if not mount_checker(str(source)):
        raise RuntimeError("コピー元がマウントポイントではありません: " + str(source))
    try:
        next(source.iterdir())
    except StopIteration:
        raise RuntimeError("コピー元が空です: " + str(source))


def run(
    config: Config,
    mount_checker: Callable[[str], bool] = os.path.ismount,
    copy_function: Callable[[Path, Path], None] = atomic_copy,
) -> Summary:
    validate_source(config.source, mount_checker)
    summary = Summary()
    lock_path = Path(str(config.database) + ".lock")
    lock_context = contextlib.nullcontext() if config.dry_run else process_lock(lock_path)
    with lock_context:
        connection = open_database(config.database, read_only=config.dry_run)
        try:
            walk_errors = []  # type: List[OSError]

            def remember_walk_error(error: OSError) -> None:
                walk_errors.append(error)
                LOGGER.error("ディレクトリを読めません: %s", error)

            for source_path in iter_target_files(config.source, remember_walk_error):
                summary.targets += 1
                try:
                    stat_result = source_path.stat()
                    relative_path = source_path.relative_to(config.source).as_posix()
                    if is_recorded(
                        connection, relative_path, stat_result.st_size, stat_result.st_mtime_ns
                    ):
                        continue
                    summary.new += 1
                    if config.initialize:
                        if not config.dry_run:
                            record_file(
                                connection,
                                relative_path,
                                stat_result.st_size,
                                stat_result.st_mtime_ns,
                                None,
                                "init",
                                "initialized",
                            )
                        LOGGER.info("初期化%s: %s", "予定" if config.dry_run else "", relative_path)
                        continue

                    captured_at, date_source = determine_datetime(source_path, stat_result)
                    day = captured_at.astimezone(JST).strftime("%Y%m%d")
                    destination_directory = config.destination / day[:4] / day
                    if config.dry_run:
                        LOGGER.info("コピー予定 [%s]: %s -> %s", date_source, source_path, destination_directory)
                        continue

                    destination_directory.mkdir(parents=True, exist_ok=True)
                    destination_path, identical, numbered = select_destination(
                        source_path, destination_directory
                    )
                    if identical:
                        summary.identical += 1
                        action = "identical"
                        LOGGER.info("同一内容のためスキップ [%s]: %s -> %s", date_source, source_path, destination_path)
                    else:
                        copy_function(source_path, destination_path)
                        summary.copied += 1
                        if numbered:
                            summary.numbered += 1
                        action = "copied"
                        LOGGER.info("コピー [%s]: %s -> %s", date_source, source_path, destination_path)
                    record_file(
                        connection,
                        relative_path,
                        stat_result.st_size,
                        stat_result.st_mtime_ns,
                        str(destination_path),
                        date_source,
                        action,
                    )
                except Exception as exc:
                    summary.failed += 1
                    LOGGER.error("処理失敗: %s: %s", source_path, exc)
            summary.failed += len(walk_errors)
        finally:
            connection.close()
    return summary


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE, help="コピー元マウントポイント")
    parser.add_argument("--destination", type=Path, default=DEFAULT_DESTINATION, help="コピー先 PhotoLibrary")
    parser.add_argument("--database", type=Path, default=DEFAULT_DATABASE, help="コピー済み記録の SQLite DB")
    parser.add_argument("--init", action="store_true", dest="initialize", help="既存対象をコピー済みとして登録")
    parser.add_argument("--dry-run", action="store_true", help="コピー・DB 登録をせず予定だけを表示")
    parser.add_argument("--verbose", action="store_true", help="解析失敗の詳細も表示")
    return parser


def main(arguments: Optional[Sequence[str]] = None) -> int:
    parser = build_parser()
    namespace = parser.parse_args(arguments)
    logging.basicConfig(
        level=logging.DEBUG if namespace.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    config = Config(
        source=namespace.source,
        destination=namespace.destination,
        database=namespace.database,
        initialize=namespace.initialize,
        dry_run=namespace.dry_run,
    )
    try:
        summary = run(config)
    except Exception as exc:
        LOGGER.error("開始できません: %s", exc)
        return 2
    print(
        "まとめ: 対象={0.targets} 新規={0.new} コピー={0.copied} "
        "同一スキップ={0.identical} 連番={0.numbered} 失敗={0.failed}".format(summary)
    )
    return 1 if summary.failed else 0


if __name__ == "__main__":
    sys.exit(main())