BoltClam/src/boltclam/winlock.py

"""Best-effort file locking for the real-time download gate.

Windows has no user-mode API that lets an ordinary process intercept every
``CreateFile`` (that is a kernel minifilter). What this module *can* do is
hold the download file open with ``dwShareMode = 0`` while ClamAV scans a
copy, which makes Windows refuse other opens for that period. If the file
was already open elsewhere, we fall back to scanning a copy and say so in
the event detail. This is a meaningful speed bump, not a guarantee against
a kernel-level bypass.
"""

from __future__ import annotations

import ctypes
import logging
import os
import shutil
from ctypes import wintypes
from pathlib import Path

logger = logging.getLogger(__name__)


class LockError(Exception):
    pass


if os.name == "nt":
    _kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
    _CreateFileW = _kernel32.CreateFileW
    _CreateFileW.argtypes = [wintypes.LPCWSTR, wintypes.DWORD, wintypes.DWORD,
                             wintypes.LPVOID, wintypes.DWORD, wintypes.DWORD,
                             wintypes.HANDLE]
    _CreateFileW.restype = wintypes.HANDLE
    _ReadFile = _kernel32.ReadFile
    _ReadFile.argtypes = [wintypes.HANDLE, wintypes.LPVOID, wintypes.DWORD,
                          ctypes.POINTER(wintypes.DWORD), wintypes.LPVOID]
    _ReadFile.restype = wintypes.BOOL
    _CloseHandle = _kernel32.CloseHandle
    _CloseHandle.argtypes = [wintypes.HANDLE]
    _CloseHandle.restype = wintypes.BOOL
    _SetFileInformationByHandle = _kernel32.SetFileInformationByHandle
    _SetFileInformationByHandle.argtypes = [wintypes.HANDLE, ctypes.c_int,
                                            wintypes.LPVOID, wintypes.DWORD]
    _SetFileInformationByHandle.restype = wintypes.BOOL

    _GENERIC_READ = 0x80000000
    _DELETE = 0x00010000
    _OPEN_EXISTING = 3
    _FILE_ATTRIBUTE_NORMAL = 0x80
    _FILE_DISPOSITION_INFO_CLASS = 4
    _INVALID_HANDLE_VALUE = ctypes.c_void_p(-1).value

    class _FILE_DISPOSITION_INFO(ctypes.Structure):
        _fields_ = [("DeleteFile", wintypes.BOOL)]
else:
    _kernel32 = None
    _FILE_DISPOSITION_INFO = None


class ExclusiveFile:
    """Open a file (share-deny on Windows) and copy it while held."""

    def __init__(self, path: Path) -> None:
        self.path = Path(path)
        self._handle = None
        self._file = None
        self._delete_access = False
        self._delete_on_close = False
        self._open = False

    @property
    def exclusive(self) -> bool:
        return bool(self._open and os.name == "nt")

    def open(self) -> None:
        if self._open:
            return
        if os.name == "nt":
            handle = _CreateFileW(
                str(self.path), _GENERIC_READ | _DELETE, 0, None, _OPEN_EXISTING,
                _FILE_ATTRIBUTE_NORMAL, None,
            )
            if handle == _INVALID_HANDLE_VALUE:
                handle = _CreateFileW(
                    str(self.path), _GENERIC_READ, 0, None, _OPEN_EXISTING,
                    _FILE_ATTRIBUTE_NORMAL, None,
                )
                self._delete_access = False
            else:
                self._delete_access = True
            if handle == _INVALID_HANDLE_VALUE:
                error = ctypes.get_last_error()
                raise LockError(f"Could not open {self.path.name} exclusively (error {error}).")
            self._handle = handle
        else:
            try:
                self._file = open(self.path, "rb")
            except OSError as exc:
                raise LockError(f"Could not open {self.path.name}: {exc}") from exc
        self._open = True

    def copy_to(self, destination: Path) -> None:
        if not self._open:
            self.open()
        destination = Path(destination)
        destination.parent.mkdir(parents=True, exist_ok=True)
        try:
            if os.name == "nt":
                buffer = ctypes.create_string_buffer(1024 * 1024)
                read = wintypes.DWORD()
                with open(destination, "wb") as output:
                    while True:
                        ok = _ReadFile(self._handle, buffer, len(buffer),
                                       ctypes.byref(read), None)
                        if not ok:
                            raise LockError(
                                f"Could not read {self.path.name} (error "
                                f"{ctypes.get_last_error()}).")
                        if read.value == 0:
                            break
                        output.write(buffer.raw[:read.value])
            else:
                with open(destination, "wb") as output:
                    shutil.copyfileobj(self._file, output, length=1024 * 1024)
        except OSError as exc:
            raise LockError(f"Could not copy {self.path.name}: {exc}") from exc

    def mark_delete_on_close(self) -> bool:
        if os.name != "nt" or not self._delete_access or self._handle is None:
            return False
        info = _FILE_DISPOSITION_INFO(True)
        ok = _SetFileInformationByHandle(
            self._handle, _FILE_DISPOSITION_INFO_CLASS, ctypes.byref(info),
            ctypes.sizeof(info),
        )
        if not ok:
            logger.warning("Could not mark %s delete-on-close (error %s)",
                           self.path, ctypes.get_last_error())
            return False
        self._delete_on_close = True
        return True

    def close(self) -> None:
        if not self._open:
            return
        if os.name == "nt":
            if self._handle is not None:
                _CloseHandle(self._handle)
                self._handle = None
        else:
            if self._file is not None:
                self._file.close()
                self._file = None
        self._open = False

    def __enter__(self) -> "ExclusiveFile":
        self.open()
        return self

    def __exit__(self, exc_type, exc, traceback) -> None:
        self.close()


def copy_for_scan(path: Path, destination: Path, attempts: int = 8,
                  delay: float = 0.4) -> tuple[ExclusiveFile, bool]:
    """Copy *path* while holding it where possible.

    Returns ``(locker, exclusive)``. The caller must keep ``locker`` open
    until the scan finishes, then close it. If the file cannot be opened
    exclusively (another process is still writing or reading it), this
    retries briefly and then copies without the lock.
    """
    import time

    last_error: Exception | None = None
    for _ in range(max(1, attempts)):
        locker = ExclusiveFile(path)
        try:
            locker.open()
            locker.copy_to(destination)
            return locker, locker.exclusive
        except LockError as exc:
            last_error = exc
            try:
                locker.close()
            except Exception:
                pass
            time.sleep(delay)
    # Last resort: a plain copy. The watcher already waited for the size to
    # stop changing, so this is still normally a complete file.
    try:
        shutil.copy2(path, destination)
    except OSError as exc:
        raise LockError(f"Could not copy {path.name}: {exc}") from exc
    logger.warning("Real-time scan could not exclusively lock %s: %s", path, last_error)
    return ExclusiveFile(path), False