bandsaunter/bandsaunter/protocols.py

"""Data protocols that carry their own framing, and so can be read outright.

The generic decoder in :mod:`bandsaunter.decode` recovers bits and has to
argue about whether they are real.  The protocols here do not have that
problem: each has a sync word to find and a checksum to verify, so a frame
either passes or it does not, and one that passes is not a guess.

Two are implemented, because between them they cover most of what an ordinary
receiver actually hears carrying words rather than measurements:

``POCSAG``
    Paging.  Still in daily use by hospitals, fire services and industrial
    plant long after the consumer pagers went away, and it carries plain
    text.  512, 1200 or 2400 baud two-level FSK.

``AX.25 / APRS``
    Amateur packet.  1200 baud AFSK on VHF, and the payload is position
    reports, weather and messages -- with the sender's callsign in the
    header, which the map already knows what to do with.
"""

from __future__ import annotations

import math

import numpy as np

from .decode import DataDecode, PulseTrain, nrz_bits, slice_fsk

__all__ = ["decode_pocsag", "decode_ax25", "PROTOCOLS", "AX25Frame",
           "POCSAG_BAUDS", "SYNC_WORD", "MAX_ADDRESS", "pocsag_bits",
           "pocsag_codeword"]


# ---------------------------------------------------------------------------
# Shared
# ---------------------------------------------------------------------------

def _find(bits: str, pattern: str, start: int = 0) -> int:
    return bits.find(pattern, start)


def _invert(bits: str) -> str:
    return bits.translate(str.maketrans("01", "10"))


# ---------------------------------------------------------------------------
# POCSAG
# ---------------------------------------------------------------------------

# The frame synchronisation codeword, sent before each batch of sixteen.
SYNC_WORD = 0x7CD215D8
SYNC_BITS = f"{SYNC_WORD:032b}"

# The three rates in the standard.  Which one a transmitter uses is not
# announced anywhere in the signal, so all three are tried and the one whose
# sync word appears is the right one -- a 32-bit pattern does not turn up in
# the wrong reading by accident.
POCSAG_BAUDS: tuple[float, ...] = (1200.0, 512.0, 2400.0)

# BCH(31,21) with an even parity bit, which is what protects each codeword.
_BCH_POLY = 0b11101101001          # x^10 + x^9 + x^8 + x^6 + x^5 + x^3 + 1

_IDLE = 0x7A89C197

# The 20 data bits of a message codeword are packed end to end and then read
# out either as 7-bit ASCII or as 4-bit digits, depending on the pager.
_NUMERIC = "0123456789*U -)("


def _bch_syndrome(word: int) -> int:
    """Zero when the codeword's BCH check passes."""
    remainder = word >> 1                     # drop the parity bit
    for shift in range(30, 9, -1):
        if remainder & (1 << shift):
            remainder ^= _BCH_POLY << (shift - 10)
    return remainder & 0x3FF


def _parity_ok(word: int) -> bool:
    return bin(word).count("1") % 2 == 0


def _correct(word: int) -> tuple[int, bool]:
    """Check a codeword, correcting a single bit error if there is one.

    Worth doing rather than discarding: one bad bit in thirty-two is exactly
    what a fading paging signal delivers, and the BCH code was put there to
    survive it.
    """
    if _bch_syndrome(word) == 0 and _parity_ok(word):
        return word, True
    for bit in range(32):
        trial = word ^ (1 << bit)
        if _bch_syndrome(trial) == 0 and _parity_ok(trial):
            return trial, True
    return word, False


def _pocsag_text(payload: str) -> tuple[str, str]:
    """Read a run of message bits as text and as digits.

    Both, because nothing in the message says which it is: an alphanumeric
    pager sends 7-bit ASCII least significant bit first, a numeric one sends
    4-bit digits, and the only way to tell is to look at what comes out.
    """
    letters = []
    for i in range(0, len(payload) - 6, 7):
        chunk = payload[i:i + 7]
        code = int(chunk[::-1], 2)            # LSB first on the air
        letters.append(chr(code) if 32 <= code < 127 else
                       ("\n" if code in (10, 13) else "."))
    digits = []
    for i in range(0, len(payload) - 3, 4):
        code = int(payload[i:i + 4][::-1], 2)
        digits.append(_NUMERIC[code])
    return "".join(letters).rstrip(". \n"), "".join(digits).rstrip(" ")


def _readable(text: str) -> float:
    """What share of a string is characters a message would really contain."""
    if not text:
        return 0.0
    good = sum(1 for c in text if c.isalnum() or c in " .,:;/-+()@#'\"!?\n")
    return good / len(text)


def _pocsag_batches(bits: str) -> list[tuple[int, str, str]]:
    """Every address and message in a bit stream.

    Returns ``(address, function, text)``.  Frames are read from each sync
    word independently, so a stream that loses lock partway through still
    yields everything before and after it.
    """
    out: list[tuple[int, str, str]] = []
    pos = 0
    pending_address: int | None = None
    pending_function = 0
    pending_bits: list[str] = []

    def flush():
        if pending_address is None:
            return
        payload = "".join(pending_bits)
        text, digits = _pocsag_text(payload)
        # Whichever reading looks more like a message someone would send.
        chosen = text if _readable(text) >= 0.75 and len(text) >= 2 else digits
        # Function bits 00..11 are the pager's four addresses, written A to
        # D by everyone who documents them.
        out.append((pending_address, "ABCD"[pending_function & 3],
                    chosen.strip()))

    while True:
        at = _find(bits, SYNC_BITS, pos)
        if at < 0:
            break
        pos = at + 32
        for frame in range(8):
            for half in range(2):
                start = pos + (frame * 2 + half) * 32
                if start + 32 > len(bits):
                    pos = len(bits)
                    break
                word = int(bits[start:start + 32], 2)
                if word == _IDLE:
                    flush()
                    pending_address, pending_bits = None, []
                    continue
                word, valid = _correct(word)
                if not valid:
                    continue
                if word & 0x80000000:
                    # A message codeword: twenty bits of payload.
                    if pending_address is not None:
                        pending_bits.append(f"{word:032b}"[1:21])
                    continue
                flush()
                # An address codeword.  The low three bits of the address are
                # not sent: they are which of the eight frames it arrived in.
                pending_address = ((word >> 13) & 0x3FFFF) << 3 | frame
                pending_function = (word >> 11) & 0x3
                pending_bits = []
            else:
                continue
            break
        pos += 16 * 32
    flush()
    return out


def decode_pocsag(x: np.ndarray, sample_rate: float,
                  baud_hint: float = 0.0) -> DataDecode | None:
    """Read POCSAG paging out of a two-level FSK signal."""
    train = slice_fsk(x, sample_rate)
    if train is None or len(train) < 16:
        return None
    order = list(POCSAG_BAUDS)
    if baud_hint:
        order.sort(key=lambda b: abs(b - baud_hint))
    for baud in order:
        raw = nrz_bits(train, baud)
        if len(raw) < 64:
            continue
        for bits in (raw, _invert(raw)):
            if SYNC_BITS not in bits:
                continue
            pages = _pocsag_batches(bits)
            if not pages:
                continue
            # A capture usually spans several batches and a transmitter
            # repeats its queue, so the same page arrives more than once.
            # Reported once, in the order it first appeared.
            messages: list[str] = []
            for address, function, text in pages:
                line = f"[{address:07d}{function}]"
                line = f"{line} {text}" if text else line
                if line not in messages:
                    messages.append(line)
            syncs = bits.count(SYNC_BITS)
            return DataDecode(
                ok=True, protocol=f"POCSAG {baud:.0f}", encoding="NRZ",
                baud=baud, bits=bits[:256], repeats=syncs,
                agreement=1.0 if syncs > 1 else 0.0,
                checks=["BCH(31,21)"], messages=messages,
                confidence=min(0.97, 0.72 + 0.05 * len(pages)
                               + 0.05 * min(3, syncs)))
    return None


# ---------------------------------------------------------------------------
# AX.25 over Bell 202 AFSK -- APRS and amateur packet
# ---------------------------------------------------------------------------

MARK_HZ = 1200.0
SPACE_HZ = 2200.0
AFSK_BAUD = 1200.0
_FLAG = "01111110"


class AX25Frame:
    """One AX.25 frame whose frame-check sequence was correct."""

    def __init__(self, raw: bytes):
        self.raw = raw
        self.source = ""
        self.destination = ""
        self.path: list[str] = []
        self.info = ""
        self._parse()

    def _parse(self) -> None:
        addresses = []
        i = 0
        while i + 7 <= len(self.raw) and len(addresses) < 10:
            field = self.raw[i:i + 7]
            call = "".join(chr(b >> 1) for b in field[:6]).strip()
            ssid = (field[6] >> 1) & 0x0F
            addresses.append(f"{call}-{ssid}" if ssid else call)
            i += 7
            if field[6] & 0x01:               # the end-of-address bit
                break
        if len(addresses) >= 2:
            self.destination, self.source = addresses[0], addresses[1]
            self.path = addresses[2:]
        # Skip the control and protocol-identifier bytes.
        body = self.raw[i + 2:] if len(self.raw) > i + 2 else b""
        self.info = body.decode("ascii", "replace").rstrip("\r\n")

    def describe(self) -> str:
        route = ">".join([self.source or "?", self.destination or "?"]
                         + self.path)
        return f"{route}: {self.info}" if self.info else route


def _fcs(data: bytes) -> int:
    """The AX.25 frame check: CRC-16/X.25, reflected, inverted at the end."""
    crc = 0xFFFF
    for byte in data:
        crc ^= byte
        for _ in range(8):
            crc = (crc >> 1) ^ 0x8408 if crc & 1 else crc >> 1
    return crc ^ 0xFFFF


def _afsk_symbols(x: np.ndarray, sample_rate: float) -> np.ndarray | None:
    """Turn Bell 202 audio into a two-level signal, mark against space.

    A correlator rather than a discriminator: the two tones are close enough
    together, and radio audio distorted enough, that measuring which of the
    two a bit-length window contains more of works where measuring the
    instantaneous frequency does not.
    """
    x = np.asarray(x)
    if np.iscomplexobj(x):
        # FM first: AFSK is audio, and on VHF it arrives inside an FM carrier.
        phase = np.unwrap(np.angle(x))
        audio = np.diff(phase)
    else:
        audio = x.astype(np.float64)
    if audio.size < int(sample_rate / 100.0):
        return None
    audio = audio - float(audio.mean())
    window = int(round(sample_rate / AFSK_BAUD))
    if window < 4:
        return None
    n = np.arange(window)
    out = []
    for tone in (MARK_HZ, SPACE_HZ):
        arg = 2.0 * math.pi * tone * n / sample_rate
        i = np.convolve(audio, np.cos(arg)[::-1], mode="same")
        q = np.convolve(audio, np.sin(arg)[::-1], mode="same")
        out.append(np.hypot(i, q))
    return out[0] - out[1]


def _hdlc_frames(bits: str) -> list[bytes]:
    """Split a bit stream at HDLC flags and undo the bit stuffing."""
    frames: list[bytes] = []
    pos = bits.find(_FLAG)
    if pos < 0:
        return frames
    while True:
        # Flags repeat back to back between frames; skip past all of them.
        while bits.startswith(_FLAG, pos):
            pos += 8
        end = bits.find(_FLAG, pos)
        if end < 0:
            break
        body = bits[pos:end]
        pos = end
        if len(body) < 8 * 17:            # shorter than an empty AX.25 frame
            continue
        # Undo stuffing: a zero inserted after every five ones.
        out, ones = [], 0
        for bit in body:
            if ones == 5:
                ones = 0
                if bit == "0":
                    continue              # the stuffed bit
            out.append(bit)
            ones = ones + 1 if bit == "1" else 0
        clean = "".join(out)
        whole = len(clean) - len(clean) % 8
        # Least significant bit first on the air.
        frames.append(bytes(int(clean[i:i + 8][::-1], 2)
                            for i in range(0, whole, 8)))
    return frames


def decode_ax25(x: np.ndarray, sample_rate: float,
                baud_hint: float = 0.0) -> DataDecode | None:
    """Read AX.25 packet, as APRS uses it, from 1200 baud AFSK."""
    soft = _afsk_symbols(x, sample_rate)
    if soft is None:
        return None
    levels = soft > 0
    edges = np.flatnonzero(np.diff(levels.astype(np.int8))) + 1
    if edges.size < 16:
        return None
    starts = np.concatenate(([0], edges))
    ends = np.concatenate((edges, [levels.size]))
    train = PulseTrain(levels[starts], (ends - starts).astype(np.int64),
                       float(sample_rate), "fsk")
    raw = nrz_bits(train, AFSK_BAUD)
    if len(raw) < 200:
        return None

    frames: list[AX25Frame] = []
    for stream in (raw, _invert(raw)):
        # NRZI: the data is in whether the level changed, not what it is.
        decoded = ["1" if a == b else "0"
                   for a, b in zip(stream, stream[1:])]
        for frame in _hdlc_frames("".join(decoded)):
            if len(frame) < 17 or _fcs(frame[:-2]) != \
                    (frame[-1] << 8 | frame[-2]):
                continue
            frames.append(AX25Frame(frame[:-2]))
        if frames:
            break
    if not frames:
        return None
    return DataDecode(
        ok=True, protocol="AX.25 / APRS", encoding="NRZI",
        baud=AFSK_BAUD, bits=raw[:256], repeats=len(frames),
        agreement=1.0, checks=["FCS (CRC-16/X.25)"],
        messages=[f.describe() for f in frames],
        confidence=min(0.98, 0.85 + 0.04 * len(frames)))


# Tried in order; each returns None when the signal is not its own.
PROTOCOLS = (decode_pocsag, decode_ax25)


# ---------------------------------------------------------------------------
# Encoding POCSAG, for the simulator and the tests
# ---------------------------------------------------------------------------
#
# Kept next to the decoder rather than in the test helpers so the two cannot
# drift apart: a bug shared by an encoder and its decoder is invisible, and
# the way to avoid one is to have exactly one copy of the polynomial, the
# sync word and the idle word for both to use.

def pocsag_codeword(payload: int) -> int:
    """Add the BCH check bits and the parity bit to 21 bits of payload."""
    word = (payload & 0x1FFFFF) << 10
    remainder = word
    for shift in range(30, 9, -1):
        if remainder & (1 << shift):
            remainder ^= _BCH_POLY << (shift - 10)
    word = (word | (remainder & 0x3FF)) << 1
    return word | (1 if bin(word).count("1") % 2 else 0)


# Eighteen bits go out on the air and three more are implied by which frame
# the codeword arrived in, so this is the whole address space.
MAX_ADDRESS = (1 << 21) - 1


def pocsag_bits(pages, preamble: int = 600) -> str:
    """A complete POCSAG transmission: preamble, then batches of sixteen.

    Raises on an address that will not fit rather than truncating it: a
    stream built around a silently mangled address decodes to a different
    pager, which is worse than not building one.
    """
    for address, _, _ in pages:
        if not 0 <= address <= MAX_ADDRESS:
            raise ValueError(
                f"POCSAG address {address} is outside 0-{MAX_ADDRESS}")
    out = ["10" * (preamble // 2)]
    batch: list[int] = []

    def flush() -> None:
        if not batch:
            return
        while len(batch) < 16:
            batch.append(_IDLE)
        out.append(SYNC_BITS + "".join(f"{w:032b}" for w in batch[:16]))
        batch.clear()

    for address, function, text in pages:
        frame = address & 0x7
        while len(batch) < frame * 2:
            batch.append(_IDLE)
        batch.append(pocsag_codeword(((address >> 3) << 2) | (function & 3)))
        payload = "".join(f"{ord(c):07b}"[::-1] for c in text)
        payload += "0" * (-len(payload) % 20)
        for i in range(0, len(payload), 20):
            batch.append(pocsag_codeword((1 << 20) | int(payload[i:i + 20], 2)))
            if len(batch) >= 16:
                flush()
    flush()
    return "".join(out)