bandsaunter/bandsaunter/librtlsdr.py

"""Minimal ctypes binding to librtlsdr.

Only the subset of the API that the scanner needs is bound.  Kept separate from
``device.py`` so that the raw C surface stays easy to audit.
"""

from __future__ import annotations

import ctypes
import ctypes.util
import os

__all__ = ["LibRtlSdr", "RtlSdrError", "load", "is_available", "load_error"]


class RtlSdrError(RuntimeError):
    """Raised when librtlsdr reports a failure or is unusable."""


# librtlsdr returns negative errno-ish codes; these are the ones worth naming.
_ERRNO_TEXT = {
    -1: "device handle is invalid",
    -2: "device not found or already claimed",
    -3: "operation not supported by this tuner",
    -5: "USB transfer error (cable, power, or driver problem)",
    -6: "device is busy",
    -12: "out of memory",
}


def _describe(code: int) -> str:
    return _ERRNO_TEXT.get(code, f"librtlsdr error {code}")


_CANDIDATES = (
    "librtlsdr.so.2",
    "librtlsdr.so.0",
    "librtlsdr.so",
    "librtlsdr.dylib",
    "rtlsdr.dll",
    "librtlsdr.dll",
)

_lib = None
_load_error: str | None = None


def _try_load():
    """Locate and dlopen librtlsdr, returning (handle, error_message)."""
    tried = []

    env = os.environ.get("BANDSAUNTER_LIBRTLSDR")
    names = ([env] if env else []) + list(_CANDIDATES)

    found = ctypes.util.find_library("rtlsdr")
    if found:
        names.append(found)

    for name in names:
        try:
            return ctypes.CDLL(name), None
        except OSError as exc:  # pragma: no cover - platform dependent
            tried.append(f"{name}: {exc}")

    hint = (
        "librtlsdr was not found.  On Debian/Ubuntu install it with:\n"
        "    sudo apt install rtl-sdr librtlsdr0\n"
        "On Fedora:   sudo dnf install rtl-sdr\n"
        "On macOS:    brew install librtlsdr\n"
        "Or set BANDSAUNTER_LIBRTLSDR=/path/to/librtlsdr.so"
    )
    return None, hint + "\n\nTried:\n  " + "\n  ".join(tried)


# Callback librtlsdr invokes from its own thread for each filled buffer.
READ_ASYNC_CB = ctypes.CFUNCTYPE(None, ctypes.POINTER(ctypes.c_ubyte),
                                 ctypes.c_uint32, ctypes.c_void_p)


def _bind(lib):
    c = ctypes.c_int
    u32 = ctypes.c_uint32
    p = ctypes.c_void_p
    cp = ctypes.c_char_p

    sig = {
        "rtlsdr_get_device_count": ([], ctypes.c_uint32),
        "rtlsdr_get_device_name": ([u32], cp),
        "rtlsdr_get_device_usb_strings": ([u32, cp, cp, cp], c),
        "rtlsdr_get_index_by_serial": ([cp], c),
        "rtlsdr_open": ([ctypes.POINTER(p), u32], c),
        "rtlsdr_close": ([p], c),
        "rtlsdr_set_center_freq": ([p, u32], c),
        "rtlsdr_get_center_freq": ([p], u32),
        "rtlsdr_set_freq_correction": ([p, c], c),
        "rtlsdr_get_freq_correction": ([p], c),
        "rtlsdr_get_tuner_type": ([p], c),
        "rtlsdr_get_tuner_gains": ([p, ctypes.POINTER(c)], c),
        "rtlsdr_set_tuner_gain": ([p, c], c),
        "rtlsdr_get_tuner_gain": ([p], c),
        "rtlsdr_set_tuner_gain_mode": ([p, c], c),
        "rtlsdr_set_tuner_bandwidth": ([p, u32], c),
        "rtlsdr_set_sample_rate": ([p, u32], c),
        "rtlsdr_get_sample_rate": ([p], u32),
        "rtlsdr_set_agc_mode": ([p, c], c),
        "rtlsdr_set_direct_sampling": ([p, c], c),
        "rtlsdr_get_direct_sampling": ([p], c),
        "rtlsdr_set_offset_tuning": ([p, c], c),
        "rtlsdr_get_offset_tuning": ([p], c),
        "rtlsdr_set_bias_tee": ([p, c], c),
        "rtlsdr_reset_buffer": ([p], c),
        "rtlsdr_read_sync": ([p, p, c, ctypes.POINTER(c)], c),
        "rtlsdr_read_async": ([p, READ_ASYNC_CB, p, u32, u32], c),
        "rtlsdr_cancel_async": ([p], c),
    }

    missing = []
    for name, (argtypes, restype) in sig.items():
        fn = getattr(lib, name, None)
        if fn is None:
            missing.append(name)
            continue
        fn.argtypes = argtypes
        fn.restype = restype

    # A few entry points only exist in newer librtlsdr; the wrapper degrades
    # gracefully for those rather than refusing to run.
    hard_required = {
        "rtlsdr_open",
        "rtlsdr_close",
        "rtlsdr_set_center_freq",
        "rtlsdr_set_sample_rate",
        "rtlsdr_read_sync",
        "rtlsdr_reset_buffer",
    }
    fatal = hard_required.intersection(missing)
    if fatal:
        raise RtlSdrError(
            "librtlsdr is missing required symbols: " + ", ".join(sorted(fatal))
        )
    return lib


def load():
    """Return the loaded librtlsdr handle, raising RtlSdrError if unavailable."""
    global _lib, _load_error
    if _lib is not None:
        return _lib
    if _load_error is not None:
        raise RtlSdrError(_load_error)
    lib, err = _try_load()
    if lib is None:
        _load_error = err
        raise RtlSdrError(err)
    _lib = _bind(lib)
    return _lib


def is_available() -> bool:
    try:
        load()
        return True
    except RtlSdrError:
        return False


def load_error() -> str | None:
    """The dlopen failure message, or None if the library loaded fine."""
    try:
        load()
        return None
    except RtlSdrError as exc:
        return str(exc)


class LibRtlSdr:
    """Thin, checked wrapper around the C calls.

    Every method raises :class:`RtlSdrError` on a negative return code so the
    layers above can assume success.
    """

    TUNER_NAMES = {
        0: "unknown",
        1: "Elonics E4000",
        2: "Fitipower FC0012",
        3: "Fitipower FC0013",
        4: "FCI FC2580",
        5: "Rafael Micro R820T/R820T2",
        6: "Rafael Micro R828D",
    }

    def __init__(self):
        self.lib = load()

    # -- enumeration ---------------------------------------------------
    def device_count(self) -> int:
        return int(self.lib.rtlsdr_get_device_count())

    def device_name(self, index: int) -> str:
        name = self.lib.rtlsdr_get_device_name(index)
        return name.decode("utf-8", "replace") if name else "?"

    def usb_strings(self, index: int):
        buf = [ctypes.create_string_buffer(256) for _ in range(3)]
        rc = self.lib.rtlsdr_get_device_usb_strings(index, *buf)
        if rc < 0:
            return ("?", "?", "?")
        return tuple(b.value.decode("utf-8", "replace") for b in buf)

    # -- lifecycle -----------------------------------------------------
    def open(self, index: int) -> ctypes.c_void_p:
        handle = ctypes.c_void_p()
        rc = self.lib.rtlsdr_open(ctypes.byref(handle), index)
        if rc < 0:
            extra = ""
            if rc == -2 or rc == -6:
                extra = (
                    "\nThe DVB-T kernel driver may have claimed it.  Blacklist it:\n"
                    "    echo 'blacklist dvb_usb_rtl28xxu' | "
                    "sudo tee /etc/modprobe.d/blacklist-rtl.conf\n"
                    "    sudo rmmod dvb_usb_rtl28xxu\n"
                    "Also make sure your user can access the USB device "
                    "(udev rule / plugdev group)."
                )
            raise RtlSdrError(f"could not open device {index}: {_describe(rc)}{extra}")
        return handle

    def close(self, dev) -> None:
        if dev:
            self.lib.rtlsdr_close(dev)

    # -- checked setters -----------------------------------------------
    def _check(self, rc: int, what: str, soft: bool = False) -> bool:
        if rc < 0:
            if soft:
                return False
            raise RtlSdrError(f"{what} failed: {_describe(rc)}")
        return True

    def set_center_freq(self, dev, hz: int) -> None:
        self._check(self.lib.rtlsdr_set_center_freq(dev, int(hz)),
                    f"tuning to {hz/1e6:.6f} MHz")

    def get_center_freq(self, dev) -> int:
        return int(self.lib.rtlsdr_get_center_freq(dev))

    def set_sample_rate(self, dev, hz: int) -> None:
        self._check(self.lib.rtlsdr_set_sample_rate(dev, int(hz)),
                    f"setting sample rate {hz}")

    def get_sample_rate(self, dev) -> int:
        return int(self.lib.rtlsdr_get_sample_rate(dev))

    def set_freq_correction(self, dev, ppm: int) -> bool:
        # librtlsdr returns -2 when the value is already set; that is benign.
        rc = self.lib.rtlsdr_set_freq_correction(dev, int(ppm))
        return rc >= 0 or rc == -2

    def get_tuner_type(self, dev) -> str:
        try:
            return self.TUNER_NAMES.get(int(self.lib.rtlsdr_get_tuner_type(dev)),
                                        "unknown")
        except Exception:
            return "unknown"

    def get_tuner_gains(self, dev) -> list[float]:
        n = self.lib.rtlsdr_get_tuner_gains(dev, None)
        if n <= 0:
            return []
        arr = (ctypes.c_int * n)()
        self.lib.rtlsdr_get_tuner_gains(dev, arr)
        return [v / 10.0 for v in arr]

    def set_tuner_gain_mode(self, dev, manual: bool) -> None:
        self._check(self.lib.rtlsdr_set_tuner_gain_mode(dev, 1 if manual else 0),
                    "setting tuner gain mode")

    def set_tuner_gain(self, dev, db: float) -> None:
        self._check(self.lib.rtlsdr_set_tuner_gain(dev, int(round(db * 10))),
                    f"setting tuner gain {db} dB")

    def get_tuner_gain(self, dev) -> float:
        return self.lib.rtlsdr_get_tuner_gain(dev) / 10.0

    def set_tuner_bandwidth(self, dev, hz: int) -> bool:
        fn = getattr(self.lib, "rtlsdr_set_tuner_bandwidth", None)
        if fn is None:
            return False
        return self._check(fn(dev, int(hz)), "setting tuner bandwidth", soft=True)

    def set_agc_mode(self, dev, on: bool) -> bool:
        return self._check(self.lib.rtlsdr_set_agc_mode(dev, 1 if on else 0),
                           "setting RTL2832 AGC", soft=True)

    def set_direct_sampling(self, dev, mode: int) -> bool:
        """0 = off, 1 = I branch, 2 = Q branch."""
        return self._check(self.lib.rtlsdr_set_direct_sampling(dev, int(mode)),
                           "setting direct sampling", soft=True)

    def get_direct_sampling(self, dev) -> int:
        try:
            return int(self.lib.rtlsdr_get_direct_sampling(dev))
        except Exception:
            return 0

    def set_offset_tuning(self, dev, on: bool) -> bool:
        return self._check(self.lib.rtlsdr_set_offset_tuning(dev, 1 if on else 0),
                           "setting offset tuning", soft=True)

    def set_bias_tee(self, dev, on: bool) -> bool:
        fn = getattr(self.lib, "rtlsdr_set_bias_tee", None)
        if fn is None:
            return False
        return self._check(fn(dev, 1 if on else 0), "setting bias tee", soft=True)

    def reset_buffer(self, dev) -> None:
        self._check(self.lib.rtlsdr_reset_buffer(dev), "resetting USB buffer")

    def read_async(self, dev, callback, buf_num: int = 15,
                   buf_len: int = 262144) -> None:
        """Stream continuously until :meth:`cancel_async`.  Blocks the caller."""
        rc = self.lib.rtlsdr_read_async(dev, callback, None,
                                        int(buf_num), int(buf_len))
        if rc < 0:
            raise RtlSdrError(f"async read failed: {_describe(rc)}")

    def cancel_async(self, dev) -> None:
        self.lib.rtlsdr_cancel_async(dev)

    def read_sync(self, dev, buf, count: int) -> int:
        n_read = ctypes.c_int(0)
        rc = self.lib.rtlsdr_read_sync(
            dev, ctypes.cast(buf, ctypes.c_void_p), count, ctypes.byref(n_read)
        )
        if rc < 0:
            raise RtlSdrError(f"USB read failed: {_describe(rc)}")
        return int(n_read.value)