functions.py

# -*- coding: utf-8 -*-

# Read more about process_vm_(read/write)v here:
# https://man7.org/linux/man-pages/man2/process_vm_readv.2.html

# Read more about proc and memory mapping here:
# https://man7.org/linux/man-pages/man5/proc.5.html

import ctypes
import errno as errno_mod
import logging
import os
from ctypes import addressof, sizeof
from typing import Dict, Generator, Optional, Sequence, Tuple, Type, TypeVar, Union

from ..enums import ScanTypesEnum
from ..process.errors import ProcessIDNotExistsError
from ..process.module_info import ModuleInfo
from ..process.region import (
    MemoryRegion,
    default_address_filter,
    default_scan_filter,
    make_region,
)
from ..process.scanning import (
    iter_pattern_results,
    iter_search_results,
    iter_values_for_addresses,
)
from ..process.thread_info import ThreadInfo
from ..util import (
    _validate_pytype,
    as_writable_c_buffer,
    get_c_type_of,
    values_to_bytes,
)
from ..util.pattern import PatternLike, compile_pattern
from .libc import libc
from .types import MEMORY_BASIC_INFORMATION, PATH_SIZE, PRIVILEGES_SIZE, iovec


_logger = logging.getLogger("PyMemoryEditor")


T = TypeVar("T")


# Errors that mean "the page is no longer mapped" — safe to skip during scans.
# Other errors (EACCES, EPERM, ESRCH, EINVAL) reveal a real problem and are
# propagated so callers can act on them.
_PAGE_GONE_ERRNOS = frozenset((errno_mod.EFAULT, errno_mod.ENOMEM))


class _LinuxPartialIOError(OSError):
    """
    process_vm_readv / process_vm_writev returned fewer bytes than requested.

    In practice this happens when the target range straddles a freed or
    inaccessible page — the kernel transfers what it can and reports the
    short count. The previous behavior was to silently accept the short
    result, leaving the caller's buffer half-filled with real bytes and
    half with zeros (which downstream decoding would treat as valid).
    Mirrors the partial-read/write check the Win32 backend already does
    against ``ReadProcessMemory`` / ``WriteProcessMemory``.
    """

    def __init__(self, op: str, address: int, bytes_done: int, length: int):
        super().__init__(
            "%s partial transfer at 0x%X: %d of %d bytes."
            % (op, address, bytes_done, length)
        )
        self.address = address
        self.bytes_done = bytes_done
        self.length = length


def _process_vm_readv(
    pid: int, local_address: int, remote_address: int, length: int
) -> int:
    """
    Wrapper for process_vm_readv that raises OSError on failure.
    Returns the number of bytes read.

    Raises ``_LinuxPartialIOError`` when the kernel reports a short read
    (``result < length``) so callers don't decode a buffer that is part
    real-bytes, part zero-initialized. Scan loops classify this as a
    transient failure (same shape as a vanished page).
    """
    local = (iovec * 1)(iovec(local_address, length))
    remote = (iovec * 1)(iovec(remote_address, length))
    result = libc.process_vm_readv(pid, local, 1, remote, 1, 0)

    if result == -1:
        errno = ctypes.get_errno()
        raise OSError(errno, os.strerror(errno))

    if result != length:
        raise _LinuxPartialIOError(
            "process_vm_readv", remote_address, result, length
        )

    return result


def _process_vm_writev(
    pid: int, local_address: int, remote_address: int, length: int
) -> int:
    """
    Wrapper for process_vm_writev that raises OSError on failure.
    Returns the number of bytes written.

    Raises ``_LinuxPartialIOError`` on a short write so the caller learns
    that the value did not fully land. The Win32 backend already enforces
    this for ``WriteProcessMemory``.
    """
    local = (iovec * 1)(iovec(local_address, length))
    remote = (iovec * 1)(iovec(remote_address, length))
    result = libc.process_vm_writev(pid, local, 1, remote, 1, 0)

    if result == -1:
        errno = ctypes.get_errno()
        raise OSError(errno, os.strerror(errno))

    if result != length:
        raise _LinuxPartialIOError(
            "process_vm_writev", remote_address, result, length
        )

    return result


def _make_read_chunk(pid: int):
    """
    Build a `read_chunk(address, size)` closure bound to ``pid``.

    Each backend used to define this closure inline three times
    (``search_addresses_by_value``, ``search_addresses_by_pattern``,
    ``search_values_by_addresses``) — drift between them was a recurring source
    of subtle bugs. Owning it once per backend keeps the trio in lockstep.
    """
    def read_chunk(address: int, size: int):
        buffer = (ctypes.c_byte * size)()
        _process_vm_readv(pid, addressof(buffer), address, sizeof(buffer))
        return buffer

    return read_chunk


def _is_transient(exc: BaseException) -> bool:
    """
    Classify ``exc`` as a transient page-vanished failure that scan loops
    should swallow vs a real configuration / permission error that must
    propagate. Shared by every scanning entry point.
    """
    # A short read mid-scan is equivalent to a page disappearing — the scan
    # should skip the chunk and keep going. Real permission/configuration
    # errors (EACCES, EPERM, ESRCH, EINVAL) propagate.
    if isinstance(exc, _LinuxPartialIOError):
        return True
    return isinstance(exc, OSError) and exc.errno in _PAGE_GONE_ERRNOS


def get_memory_regions(pid: int) -> Generator["MemoryRegion", None, None]:
    """
    Yield a :class:`MemoryRegion` for each entry in ``/proc/<pid>/maps``.

    Translates the typical I/O failures of ``/proc/<pid>/maps`` into the
    library's own exception hierarchy so callers don't have to special-case
    raw ``FileNotFoundError`` / ``PermissionError`` from the kernel pseudo-fs.
    """
    mapping_filename = "/proc/{}/maps".format(pid)

    try:
        mapping_file = open(mapping_filename, "r")
    except FileNotFoundError:
        # Target died between OpenProcess()'s pid_exists() check and now. The
        # caller already accepts ProcessIDNotExistsError from the open path,
        # so funnel into it instead of leaking a kernel pseudo-fs error.
        raise ProcessIDNotExistsError(pid)
    except PermissionError as exc:
        # ptrace_scope (or being a non-root user inspecting a different uid)
        # is the typical reason — re-raise with a hint pointing at the fix.
        raise PermissionError(
            "Cannot read %s: %s. On Linux this usually means ptrace_scope "
            "is restricting access; try `sudo sysctl kernel.yama.ptrace_scope=0` "
            "or run as root." % (mapping_filename, exc)
        )

    with mapping_file:
        for line in mapping_file:
            region_information = line.split(maxsplit=5)

            try:
                addressing_range, privileges, offset, device, inode = (
                    region_information[0:5]
                )
                path = region_information[5].strip() if len(region_information) >= 6 else ""

                start_address, end_address = [
                    int(addr, 16) for addr in addressing_range.split("-")
                ]
                major_id, minor_id = [int(_id, 16) for _id in device.split(":")]

                offset = int(offset, 16)
                inode = int(inode)  # /proc/<pid>/maps formats the inode as decimal.
            except (ValueError, IndexError) as exc:
                # A single malformed line (kernel quirk, racing teardown) must
                # not abort the whole region walk — skip it and keep going, the
                # same log-and-continue contract the Windows/macOS walkers use.
                _logger.debug(
                    "get_memory_regions: skipping unparseable maps line %r: %s",
                    line,
                    exc,
                )
                continue

            size = end_address - start_address

            # Truncate to fit the fixed-size inline byte arrays in the struct.
            # Leave room for a null so attribute reads always terminate cleanly.
            privileges_bytes = privileges.encode()[: PRIVILEGES_SIZE - 1]
            path_bytes = path.encode()[: PATH_SIZE - 1]

            region = MEMORY_BASIC_INFORMATION(
                start_address,
                size,
                privileges_bytes,
                offset,
                major_id,
                minor_id,
                inode,
                path_bytes,
            )
            yield make_region(
                address=start_address,
                size=region.RegionSize,
                struct=region,
            )


def get_modules(pid: int) -> Generator[ModuleInfo, None, None]:
    """
    Yield a :class:`ModuleInfo` for every file-backed module mapped by the
    process, derived from ``/proc/<pid>/maps``.

    A single library spans several consecutive mappings (its ``.text``,
    ``.rodata``, ``.data`` / ``.bss`` segments), all sharing the same backing
    path. Mappings are grouped by path: the module's ``base_address`` is the
    lowest mapping start and its ``size`` reaches the highest mapping end.

    Pseudo-mappings without a real backing file (``[heap]``, ``[stack]``,
    ``[vvar]``, anonymous memory) are not modules and are skipped — only paths
    beginning with ``/`` are considered. Modules are yielded in ascending
    ``base_address`` order, which is the order ``/proc/<pid>/maps`` lists them.
    """
    # path -> [min_start, max_end]; first_seen preserves the maps order, which
    # is already ascending by address, so we don't need to sort afterwards.
    bounds: Dict[str, list] = {}
    first_seen = []

    for region in get_memory_regions(pid):
        path = region.path
        if not path.startswith("/"):
            continue

        start = region.address
        end = start + region.size

        entry = bounds.get(path)
        if entry is None:
            bounds[path] = [start, end]
            first_seen.append(path)
        else:
            if start < entry[0]:
                entry[0] = start
            if end > entry[1]:
                entry[1] = end

    for path in first_seen:
        start, end = bounds[path]
        yield ModuleInfo(
            name=os.path.basename(path),
            path=path,
            base_address=start,
            size=end - start,
            raw=path,
        )


def _read_elf_class(path: str) -> Optional[int]:
    """
    Read the ``EI_CLASS`` byte from an ELF file's identification header.

    Returns 1 for ELFCLASS32, 2 for ELFCLASS64, or ``None`` if ``path`` can't
    be opened or isn't an ELF file (missing ``\\x7fELF`` magic).
    """
    try:
        with open(path, "rb") as elf:
            ident = elf.read(5)
    except OSError:
        return None

    if len(ident) < 5 or ident[:4] != b"\x7fELF":
        return None

    return ident[4]  # e_ident[EI_CLASS]: 1 = ELFCLASS32, 2 = ELFCLASS64


def _detect_process_64bit(pid: int) -> Optional[bool]:
    """
    Return ``True``/``False`` from the ELF ``EI_CLASS`` byte of the target's
    executable, or ``None`` when no header could be read. The raw *mechanism*:
    no guessing and no warning — the caller decides what an unknown result
    means (the public :func:`is_process_64bit` falls back to the host word size;
    ``AbstractProcess.is_64bit`` honors ``strict_bitness``).

    The primary source is ``/proc/<pid>/exe``; if that symlink can't be read (a
    different user without ``CAP_SYS_PTRACE``), it falls back to the first
    file-backed, executable mapping in ``/proc/<pid>/maps`` — the main image or
    a shared library, which share the process's bitness.
    """
    ei_class = _read_elf_class("/proc/{}/exe".format(pid))

    if ei_class is None:
        # Fallback: probe a file-backed executable mapping's on-disk ELF header.
        for region in get_memory_regions(pid):
            if not region.is_executable:
                continue
            path = region.path
            if not path or path.startswith("["):  # skip [heap], [stack], anon
                continue
            ei_class = _read_elf_class(path)
            if ei_class is not None:
                break

    if ei_class == 2:
        return True
    if ei_class == 1:
        return False
    return None


def is_process_64bit(pid: int) -> bool:
    """
    Return ``True`` if the target process is 64-bit, ``False`` if 32-bit.

    Thin *policy* wrapper over :func:`_detect_process_64bit`: when no ELF class
    can be read it assumes the host's word size (the usual case) and warns so a
    wrong pointer-width default (used by the pointer APIs) is traceable instead
    of a silent mis-detection on a cross-bitness target.
    """
    detected = _detect_process_64bit(pid)
    if detected is not None:
        return detected

    _logger.warning(
        "is_process_64bit: could not read the ELF class for pid %d; assuming "
        "the host word size. Pointer-width detection may be wrong for a "
        "cross-bitness target.",
        pid,
    )
    return ctypes.sizeof(ctypes.c_void_p) == 8


def read_process_memory(pid: int, address: int, pytype: Type[T], bufflength: int) -> T:
    """
    Return a value from a memory address.
    """
    _validate_pytype(pytype)

    data = get_c_type_of(pytype, bufflength)
    _process_vm_readv(pid, addressof(data), address, sizeof(data))

    if pytype is str:
        return bytes(data).decode("utf-8", errors="replace")
    elif pytype is bytes:
        return bytes(data)
    else:
        return data.value


def read_process_memory_into(pid: int, address: int, buffer) -> int:
    """
    Read ``len(buffer)`` bytes from ``address`` directly into the writable
    ``buffer``, with no intermediate allocation. Returns the number of bytes
    read (always the buffer's byte length on success; a short read raises
    ``_LinuxPartialIOError``).
    """
    c_buffer = as_writable_c_buffer(buffer)
    size = len(c_buffer)
    return _process_vm_readv(pid, addressof(c_buffer), address, size)


def search_addresses_by_value(
    pid: int,
    pytype: Type[T],
    bufflength: int,
    value: Union[bool, int, float, str, bytes, tuple],
    scan_type: ScanTypesEnum = ScanTypesEnum.EXACT_VALUE,
    progress_information: bool = False,
    writeable_only: bool = False,
    *,
    memory_regions: Optional[Sequence[MemoryRegion]] = None,
) -> Generator[Union[int, Tuple[int, dict]], None, None]:
    """
    Search the whole memory space, accessible to the process,
    for the provided value, returning the found addresses.

    Passing a `memory_regions` snapshot skips region enumeration.
    """
    _validate_pytype(pytype)

    target_value_bytes = values_to_bytes(pytype, bufflength, value)

    source_regions = (
        memory_regions if memory_regions is not None else get_memory_regions(pid)
    )

    filtered_regions = [
        region
        for region in source_regions
        if default_scan_filter(region, writeable_only=writeable_only)
    ]
    filtered_regions.sort(key=lambda region: region.address)

    yield from iter_search_results(
        filtered_regions,
        pytype,
        bufflength,
        target_value_bytes,
        scan_type,
        _make_read_chunk(pid),
        progress_information=progress_information,
        transient_error_check=_is_transient,
    )


def search_addresses_by_pattern(
    pid: int,
    pattern: PatternLike,
    *,
    byte_length: int = 0,
    progress_information: bool = False,
    memory_regions: Optional[Sequence[MemoryRegion]] = None,
) -> Generator[Union[int, Tuple[int, dict]], None, None]:
    """
    AOB scan against every readable, non-shared region of the target. See
    :meth:`AbstractProcess.search_by_pattern`.
    """
    compiled, length = compile_pattern(pattern, byte_length=byte_length)

    source_regions = (
        memory_regions if memory_regions is not None else get_memory_regions(pid)
    )

    filtered_regions = [
        region for region in source_regions if default_scan_filter(region)
    ]
    filtered_regions.sort(key=lambda region: region.address)

    yield from iter_pattern_results(
        filtered_regions,
        compiled,
        length,
        _make_read_chunk(pid),
        progress_information=progress_information,
        transient_error_check=_is_transient,
    )


def search_values_by_addresses(
    pid: int,
    pytype: Type[T],
    bufflength: int,
    addresses: Sequence[int],
    *,
    memory_regions: Optional[Sequence[MemoryRegion]] = None,
    raise_error: bool = False,
) -> Generator[Tuple[int, Optional[T]], None, None]:
    """
    Search the whole memory space, accessible to the process,
    for the provided list of addresses, returning their values.

    Memory is read in chunks (see iter_region_chunks) to bound the per-call
    allocation. Chunks near an address boundary read `bufflength - 1` extra
    bytes so values straddling the boundary are still decoded correctly.
    Addresses that fall in gaps between regions or extend past a region's end
    yield `(address, None)`.
    """
    _validate_pytype(pytype)

    # `None` means "no snapshot provided, enumerate now". An empty list passed
    # explicitly is honored verbatim — scanning nothing is a valid choice when
    # the caller pre-filtered to zero regions.
    if memory_regions is None:
        memory_regions = [
            region for region in get_memory_regions(pid) if default_address_filter(region)
        ]
    else:
        memory_regions = list(memory_regions)

    yield from iter_values_for_addresses(
        addresses,
        memory_regions,
        pytype,
        bufflength,
        _make_read_chunk(pid),
        raise_error=raise_error,
        transient_error_check=_is_transient,
    )


def write_process_memory(
    pid: int,
    address: int,
    pytype: Type[T],
    bufflength: int,
    value: Union[bool, int, float, str, bytes],
) -> Union[bool, int, float, str, bytes]:
    """
    Write a value to a memory address.
    """
    _validate_pytype(pytype)

    data = get_c_type_of(pytype, bufflength)
    data.value = value.encode() if isinstance(value, str) else value

    _process_vm_writev(pid, addressof(data), address, sizeof(data))
    return value


def get_threads(pid: int) -> Generator[ThreadInfo, None, None]:
    """
    Yield a :class:`ThreadInfo` for every thread of the target process by
    listing ``/proc/<pid>/task/`` — each subdirectory there is a TID.

    State and priority come from ``/proc/<pid>/task/<tid>/stat`` when readable;
    silent on permission/race errors (a thread may exit between listing the
    directory and reading its stat file) but logged at DEBUG so observers can
    see it.
    """
    task_dir = "/proc/{}/task".format(pid)

    try:
        entries = os.listdir(task_dir)
    except OSError as exc:
        _logger.debug("get_threads: could not list %s: %s", task_dir, exc)
        return

    for entry in entries:
        try:
            tid = int(entry)
        except ValueError:
            continue

        state: Optional[str] = None
        priority: Optional[int] = None
        try:
            with open("{}/{}/stat".format(task_dir, entry), "r") as fh:
                raw_stat = fh.read()
            # /proc/<pid>/task/<tid>/stat layout (man 5 proc):
            #   pid (comm) state ppid pgrp session tty_nr tpgid flags minflt
            #   cminflt majflt cmajflt utime stime cutime cstime priority ...
            # ``comm`` is wrapped in parens and *may itself contain whitespace
            # or parentheses*, so the only safe split point is the last ')'.
            close_paren = raw_stat.rfind(")")
            if close_paren != -1:
                rest = raw_stat[close_paren + 1 :].split()
                # rest[0] = state, rest[15] = priority (field 18 in the man
                # page, with the first two fields already consumed).
                if rest:
                    state = rest[0]
                if len(rest) > 15:
                    try:
                        priority = int(rest[15])
                    except ValueError:
                        priority = None
        except OSError as exc:
            _logger.debug(
                "get_threads: could not read stat for tid=%s: %s", entry, exc
            )

        yield ThreadInfo(
            tid=tid,
            start_address=None,
            state=state,
            priority=priority,
            raw=entry,
        )


def get_processes() -> Generator[Tuple[int, str], None, None]:
    """
    Yield ``(pid, name)`` for every process by listing ``/proc`` — each numeric
    subdirectory is a live pid.

    ``name`` comes from ``/proc/<pid>/comm`` (the kernel truncates it to 15
    characters via ``TASK_COMM_LEN``). Processes that vanish mid-scan are
    skipped silently (logged at DEBUG).
    """
    try:
        entries = os.listdir("/proc")
    except OSError as exc:
        _logger.debug("get_processes: could not list /proc: %s", exc)
        return

    for entry in entries:
        if not entry.isdigit():
            continue
        pid = int(entry)

        try:
            with open("/proc/{}/comm".format(entry), "r") as fh:
                name = fh.readline().rstrip("\n")
        except OSError as exc:
            # Race (process exited) or permission issue — skip it.
            _logger.debug("get_processes: could not read comm for pid=%s: %s", entry, exc)
            continue

        yield pid, name


def process_exists(pid: int) -> bool:
    """Return whether a process with ``pid`` currently exists (``/proc/<pid>``)."""
    if pid < 0:
        return False
    return os.path.isdir("/proc/{}".format(pid))