icloud_linux_app/icloud_photos_fs.py

#!/usr/bin/env python3
"""
iCloud Photos FUSE filesystem — uses fuse-python (same library as icloud-linux).

Directory layout mirrors icloudpd output: /YYYY/MM/DD/filename.HEIC
- Listing: instant (index from icloudpd --only-print-filenames)
- Open:    downloads on demand via icloudpd, cached in ~/.cache/icloud-photos/
"""

import os
import sys
import stat
import errno
import time
import threading
import subprocess
import logging
from pathlib import Path
from collections import defaultdict

import fuse
from fuse import Fuse

if not hasattr(fuse, "__version__"):
    fuse.__version__ = "0.2"
fuse.fuse_python_api = (0, 2)

log = logging.getLogger("icloud-photos-fs")

CACHE_DIR = Path.home() / ".cache" / "icloud-photos"
ICLOUDPD  = str(Path(__file__).parent.parent / "icloud-linux" / "icloud-linux" / ".venv" / "bin" / "icloudpd")
# Fallback paths
for _candidate in [
    Path.home() / "icloud-linux" / "icloud-linux" / ".venv" / "bin" / "icloudpd",
    Path.home() / "icloud-linux" / ".venv" / "bin" / "icloudpd",
]:
    if _candidate.exists():
        ICLOUDPD = str(_candidate)
        break


class MyStat(fuse.Stat):
    def __init__(self):
        self.st_mode  = 0
        self.st_ino   = 0
        self.st_dev   = 0
        self.st_nlink = 0
        self.st_uid   = os.getuid()
        self.st_gid   = os.getgid()
        self.st_size  = 0
        self.st_atime = 0
        self.st_mtime = 0
        self.st_ctime = 0


class FileEntry:
    __slots__ = ("rel_path", "mtime")

    def __init__(self, rel_path: str, mtime: float = 0.0):
        self.rel_path = rel_path
        self.mtime    = mtime

    @property
    def cached(self) -> Path:
        return CACHE_DIR / self.rel_path.lstrip("/")

    @property
    def size(self) -> int:
        cp = self.cached
        if cp.exists():
            return cp.stat().st_size
        # Return a large placeholder so callers don't truncate reads before download
        return 8 * 1024 * 1024  # 8 MB


class ICloudPhotosFS(Fuse):
    def __init__(self, username, password, cookie_dir, library,
                 on_progress=None, on_done=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._username   = username
        self._password   = password
        self._cookie_dir = cookie_dir
        self._library    = library
        self._on_progress = on_progress
        self._on_done     = on_done
        self._lock        = threading.RLock()
        self._mount_time  = time.time()

        # rel_path → FileEntry  (key has leading /)
        self._files: dict[str, FileEntry] = {}
        # "2024/06/12" → set of filenames
        self._dirs: dict[str, set] = defaultdict(set)

        CACHE_DIR.mkdir(parents=True, exist_ok=True)

    def fsinit(self):
        """Called after mount — start indexing, then prefetch recent photos."""
        threading.Thread(target=self._start_tasks, daemon=True).start()

    def _start_tasks(self):
        self._build_index()
        # Prefetch BEFORE signaling done so Photos sees real cached files immediately
        self._prefetch_recent(200)
        # Now tell the GUI we're ready
        if self._on_done:
            self._on_done(getattr(self, "_index_count", 0))

    def _prefetch_recent(self, count: int):
        log.info(f"Prefetching {count} most recent photos…")
        cmd = [
            ICLOUDPD,
            "--username",         self._username,
            "--password",         self._password,
            "--cookie-directory", self._cookie_dir,
            "--library",          self._library,
            "--no-progress-bar",
            "--directory",        str(CACHE_DIR),
            "--recent",           str(count),
            "--convert-all-heic-to-jpeg",
        ]
        try:
            subprocess.run(cmd, capture_output=True, timeout=600)
            log.info("Prefetch done")
        except Exception as e:
            log.error(f"Prefetch error: {e}")

    # ── Index ──────────────────────────────────────────────────────────────────

    def _build_index(self):
        log.info("Building index via icloudpd…")
        count = 0
        tmp_root = str(CACHE_DIR / "_index_root")
        Path(tmp_root).mkdir(parents=True, exist_ok=True)

        cmd = [
            ICLOUDPD,
            "--username",         self._username,
            "--password",         self._password,
            "--cookie-directory", self._cookie_dir,
            "--library",          self._library,
            "--no-progress-bar",
            "--directory",        tmp_root,
            "--only-print-filenames",
        ]

        try:
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                    stderr=subprocess.DEVNULL, text=True, bufsize=1)
            for line in proc.stdout:
                line = line.strip()
                if not line:
                    continue
                try:
                    rel = os.path.relpath(line, tmp_root).replace(os.sep, "/")
                except ValueError:
                    continue
                if rel.startswith("..") or not rel:
                    continue
                parts = rel.split("/")
                if len(parts) < 2:
                    continue

                filename = parts[-1]
                ext = os.path.splitext(filename)[1].lower()

                # Skip video sidecars — Photos app doesn't want .MOV/.MP4 here
                if ext in (".mov", ".mp4", ".m4v", ".avi"):
                    continue

                # Present HEIC as .jpg (we serve converted JPEG bytes)
                if ext in (".heic", ".heif"):
                    filename = os.path.splitext(filename)[0] + ".jpg"
                    parts = parts[:-1] + [filename]
                    rel = "/".join(parts)

                fuse_path = "/" + rel
                dir_key   = "/".join(parts[:-1])

                with self._lock:
                    if fuse_path not in self._files:
                        # Store original rel_path so _download knows the real date path
                        self._files[fuse_path] = FileEntry(fuse_path)
                        self._dirs[dir_key].add(filename)
                        for depth in range(1, len(parts)):
                            _ = self._dirs["/".join(parts[:depth])]
                        count += 1
                        if count % 200 == 0 and self._on_progress:
                            self._on_progress(count)
            proc.wait()
        except Exception as e:
            log.error(f"Index failed: {e}")

        log.info(f"Index done: {count}")
        # Don't call on_done yet — prefetch runs first in _start_tasks
        self._index_count = count

    # ── FUSE ops ───────────────────────────────────────────────────────────────

    def getattr(self, path):
        st = MyStat()
        t  = int(self._mount_time)

        if path == "/":
            st.st_mode  = stat.S_IFDIR | 0o555
            st.st_nlink = 2
            st.st_atime = st.st_mtime = st.st_ctime = t
            return st

        with self._lock:
            if path in self._files:
                entry = self._files[path]
                st.st_mode  = stat.S_IFREG | 0o444
                st.st_nlink = 1
                st.st_size  = entry.size
                st.st_atime = st.st_mtime = st.st_ctime = int(entry.mtime or t)
                return st
            dir_key = path.lstrip("/")
            if dir_key in self._dirs:
                st.st_mode  = stat.S_IFDIR | 0o555
                st.st_nlink = 2
                st.st_atime = st.st_mtime = st.st_ctime = t
                return st

        return -errno.ENOENT

    def readdir(self, path, offset):
        yield fuse.Direntry(".")
        yield fuse.Direntry("..")
        dir_key = path.lstrip("/")
        seen = set()

        with self._lock:
            # Sub-directories
            for k in list(self._dirs.keys()):
                if dir_key == "":
                    top = k.split("/")[0]
                    if top not in seen:
                        seen.add(top)
                        yield fuse.Direntry(top)
                elif k.startswith(dir_key + "/"):
                    rest  = k[len(dir_key) + 1:]
                    child = rest.split("/")[0]
                    if child not in seen:
                        seen.add(child)
                        yield fuse.Direntry(child)
            # Files
            for fp in list(self._files.keys()):
                fp_dir = "/".join(fp.lstrip("/").split("/")[:-1])
                if fp_dir == dir_key:
                    yield fuse.Direntry(fp.split("/")[-1])

    # Per-date download locks so concurrent reads for same day don't double-download
    _dl_locks: dict = {}
    _dl_locks_mutex = threading.Lock()
    _dl_dates_done: set = set()  # dates already fully downloaded

    def _get_dl_lock(self, key):
        with self._dl_locks_mutex:
            if key not in self._dl_locks:
                self._dl_locks[key] = threading.Lock()
            return self._dl_locks[key]

    def _ensure_downloaded(self, path: str, entry: "FileEntry"):
        """Download the file's day-batch if not already cached."""
        if entry.cached.exists():
            return
        parts = path.lstrip("/").split("/")
        try:
            date_str = f"{parts[0]}-{parts[1]}-{parts[2]}"
        except IndexError:
            date_str = None
        lock_key = date_str or path
        with self._get_dl_lock(lock_key):
            if entry.cached.exists():
                return  # another thread beat us to it
            if date_str and date_str in self._dl_dates_done:
                return  # already downloaded this day
            self._download(path, entry, date_str)
            if date_str:
                self._dl_dates_done.add(date_str)

    def open(self, path, flags):
        with self._lock:
            entry = self._files.get(path)
        if entry is None:
            return -errno.ENOENT
        write_flags = os.O_WRONLY | os.O_RDWR | os.O_CREAT | os.O_TRUNC | os.O_APPEND
        if flags & write_flags:
            return -errno.EROFS
        if not entry.cached.exists():
            log.info(f"Downloading on open: {path}")
            self._ensure_downloaded(path, entry)
        return 0

    def read(self, path, size, offset):
        with self._lock:
            entry = self._files.get(path)
        if entry is None:
            return -errno.ENOENT
        cp = entry.cached
        if not cp.exists():
            self._ensure_downloaded(path, entry)
        if not cp.exists():
            return b""
        try:
            with open(cp, "rb") as f:
                f.seek(offset)
                return f.read(size)
        except Exception:
            return b""

    def _download(self, fuse_path: str, entry: "FileEntry", date_str=None):
        entry.cached.parent.mkdir(parents=True, exist_ok=True)
        cmd = [
            ICLOUDPD,
            "--username",         self._username,
            "--password",         self._password,
            "--cookie-directory", self._cookie_dir,
            "--library",          self._library,
            "--no-progress-bar",
            "--directory",        str(CACHE_DIR),
            "--convert-all-heic-to-jpeg",
        ]
        if date_str:
            cmd += ["--skip-created-before", date_str,
                    "--skip-created-after",  date_str]
        try:
            subprocess.run(cmd, capture_output=True, timeout=120)
        except Exception as e:
            log.error(f"Download error: {e}")

    # Block writes
    def write(self, *_):   return -errno.EROFS
    def mknod(self, *_):   return -errno.EROFS
    def unlink(self, *_):  return -errno.EROFS
    def mkdir(self, *_):   return -errno.EROFS
    def rmdir(self, *_):   return -errno.EROFS
    def rename(self, *_):  return -errno.EROFS
    def symlink(self, *_): return -errno.EROFS
    def link(self, *_):    return -errno.EROFS
    def chmod(self, *_):   return -errno.EROFS
    def chown(self, *_):   return -errno.EROFS
    def truncate(self, *_):return -errno.EROFS
    def utime(self, *_):   return 0
    def statfs(self):
        st = fuse.StatVfs()
        st.f_bsize  = 4096
        st.f_blocks = 0
        st.f_bfree  = 0
        st.f_bavail = 0
        return st


def mount(username, password, cookie_dir, mountpoint, library="PrimarySync",
          on_progress=None, on_done=None):
    if os.path.ismount(mountpoint):
        # -z = lazy unmount (detaches even if busy)
        os.system(f"fusermount -uz {mountpoint!r}")
        time.sleep(1)
    os.makedirs(mountpoint, exist_ok=True)

    server = ICloudPhotosFS(
        username, password, cookie_dir, library,
        on_progress=on_progress, on_done=on_done,
        version="%prog 0.1",
        usage="iCloud Photos FUSE",
        dash_s_do="setsingle",
    )
    server.parse(args=[mountpoint, "-f"], errex=1)
    server.main()


def unmount(mountpoint: str):
    os.system(f"fusermount -u {mountpoint!r}")