pm3py/pm3py/lf/fsk.py

"""FSK demodulation — a faithful Python port of the Proxmark FSK stack.

Ported function-for-function from ``firmware/common/lfdemod.c`` (``fsk_wave_demod`` +
``aggregate_bits`` = ``fskdemod``, ``preambleSearch``, ``HIDdemodFSK``, ``detectAWID``) so it
tracks the C reference — and therefore real-tag output — rather than being a fresh reinvention.
The C thresholds against the signal mean and counts samples between 0->1 transitions: a short
wave (~fc/8) is one symbol, a long wave (~fc/10) the other; ``aggregate_bits`` then collapses
runs of like waves into clock-resolution bits. HID/AWID are FSK2a (fc/10 high, fc/8 low, RF/50).

Everything downstream (Wiegand fields, parity) is validated by the shipped ``HIDProxTag`` model.
"""
from __future__ import annotations

# lfdemod.h NOISE_AMPLITUDE_THRESHOLD — below this peak-to-mean swing the capture is "noise".
# Conservative; a real hardware capture tunes this alongside the ASK path.
NOISE_AMPLITUDE_THRESHOLD = 8


def _mean(samples) -> float:
    return sum(samples) / len(samples) if samples else 0.0


def is_noise(samples) -> bool:
    """Port of signalprop.isnoise: too little amplitude above the mean to carry data."""
    if len(samples) < 64:
        return True
    return (max(samples) - _mean(samples)) < NOISE_AMPLITUDE_THRESHOLD


def find_mod_start(src, mean: float, exp_wave: int) -> int:
    """Port of ``findModStart`` — skip leading noise/slow chip start-up to the first clean wave."""
    n = len(src)
    if n < 21:
        return 0
    i = 1
    wave_cnt = 0
    thresh_cnt = 0
    above = src[0] >= mean
    while i < n - 20:
        if src[i] < mean and above:
            thresh_cnt += 1
            if thresh_cnt > 2 and wave_cnt < exp_wave + 1:
                break
            above = False
            wave_cnt = 0
        elif src[i] >= mean and not above:
            thresh_cnt += 1
            if thresh_cnt > 2 and wave_cnt < exp_wave + 1:
                break
            above = True
            wave_cnt = 0
        else:
            wave_cnt += 1
        if thresh_cnt > 10:
            break
        i += 1
    return i


def fsk_wave_demod(samples, fchigh: int = 10, fclow: int = 8):
    """Port of ``fsk_wave_demod``: envelope -> one bit per FSK wave (1 = short/fc-low wave,
    0 = long/fc-high wave), counting samples between consecutive 0->1 transitions. Returns
    ``(bits, start_idx)``."""
    n = len(samples)
    if n < 1024:
        return [], 0
    mean = _mean(samples)
    thr = [0 if s < mean else 1 for s in samples]

    idx = find_mod_start(samples, mean, fchigh)
    last_transition = idx
    idx += 1
    out: list[int] = []
    start_idx = 0
    pre_last = last = curr = 0

    while idx < n - 20:
        if thr[idx - 1] < thr[idx]:                    # 0 -> 1 transition
            pre_last, last, curr = last, curr, idx - last_transition
            if curr < fclow - 2:
                pass                                   # garbage noise
            elif curr < fchigh - 1:                    # short wave -> 1
                if len(out) > 1 and last > fchigh - 2 and pre_last < fchigh - 1:
                    out[-1] = 1                        # fix a 9 surrounded by 8s
                out.append(1)
                if out and start_idx == 0:
                    start_idx = idx - fclow
            elif curr > fchigh + 1 and len(out) < 3:
                out = []                               # leading garbage -> reset
            elif curr == fclow + 1 and last == fclow - 1:  # 7 then 9 -> two 8s
                out.append(1)
                if out and start_idx == 0:
                    start_idx = idx - fclow
            else:                                      # long wave -> 0
                out.append(0)
                if out and start_idx == 0:
                    start_idx = idx - fchigh
            last_transition = idx
        idx += 1
    return out, start_idx


def aggregate_bits(bits, clk: int, invert: int, fchigh: int, fclow: int):
    """Port of ``aggregate_bits``: collapse runs of like waves into clock-resolution data bits
    (a run of n waves at field-clock fc spans n*fc/clk bits)."""
    if not bits:
        return []
    out: list[int] = []
    lastval = bits[0]
    n = 1
    hclk = clk // 2
    i = 1
    for i in range(1, len(bits)):
        n += 1
        if bits[i] == lastval:
            continue
        if bits[i - 1] == 1:
            n = (n * fclow + hclk) // clk
        else:
            n = (n * fchigh + hclk) // clk
        if n == 0:
            n = 1
        out.extend([bits[i - 1] ^ invert] * n)
        n = 0
        lastval = bits[i]
    if n > clk // fchigh and len(bits) >= 2:            # trailing same-frequency run
        if bits[i - 1] == 1:
            n = (n * fclow + hclk) // clk
        else:
            n = (n * fchigh + hclk) // clk
        out.extend([bits[-1] ^ invert] * n)
    return out


def fskdemod(samples, rflen: int = 50, invert: int = 1, fchigh: int = 10, fclow: int = 8):
    """Full FSK demod: envelope -> decoded 1s/0s (``fsk_wave_demod`` + ``aggregate_bits``)."""
    if is_noise(samples):
        return []
    waves, _ = fsk_wave_demod(list(samples), fchigh, fclow)
    return aggregate_bits(waves, rflen, invert, fchigh, fclow)


def preamble_search(bits, preamble):
    """Port of ``preambleSearch``: find ``preamble`` in ``bits``. Returns ``(found, start_idx,
    size)`` where ``size`` is the gap to the second occurrence (the frame length) when there is
    one, else the remaining length."""
    plen = len(preamble)
    size = len(bits)
    if size <= plen:
        return (False, 0, size)
    found = 0
    start = 0
    for idx in range(size - plen):
        if bits[idx:idx + plen] == preamble:
            found += 1
            if found == 1:
                start = idx
            elif found == 2:
                return (True, start, idx - start)
    return (found > 0, start, size)


# HID: 00011101 start-of-frame, then Manchester (10 -> 1, 01 -> 0)
HID_PREAMBLE = [0, 0, 0, 1, 1, 1, 0, 1]
# AWID: 00000001 preamble, 96-bit frame
AWID_PREAMBLE = [0, 0, 0, 0, 0, 0, 0, 1]

_M32 = 0xFFFFFFFF


def hid_demod_fsk(samples):
    """Port of ``HIDdemodFSK``: FSK2a demod + preamble + Manchester -> ``(hi2, hi, lo, nbits)``
    96-bit shift register plus the number of data bits Manchester-decoded (needed to tell the
    header-less 37-bit format from the 26-bit one), or ``None``. Field extraction is the caller's."""
    bits = fskdemod(samples, 50, 1, 10, 8)
    if len(bits) < 96 * 2:
        return None
    found, start, size = preamble_search(bits, HID_PREAMBLE)
    if not found:
        return None
    num_start = start + len(HID_PREAMBLE)
    hi2 = hi = lo = 0
    nbits = 0
    idx = num_start
    while (idx - num_start) < size - len(HID_PREAMBLE) and idx + 1 < len(bits):
        a, b = bits[idx], bits[idx + 1]
        if a == b:
            break                                      # not Manchester -> end of frame
        hi2 = ((hi2 << 1) | (hi >> 31)) & _M32
        hi = ((hi << 1) | (lo >> 31)) & _M32
        lo = ((lo << 1) | (1 if (a and not b) else 0)) & _M32
        nbits += 1
        idx += 2
    if hi == 0 and lo == 0:
        return None
    return (hi2, hi, lo, nbits)


def awid_demod_fsk(samples):
    """Port of ``detectAWID``: FSK2a demod + 00000001 preamble -> the 96 frame bits, or ``None``."""
    bits = fskdemod(samples, 50, 1, 10, 8)
    if len(bits) < 96:
        return None
    found, start, size = preamble_search(bits, AWID_PREAMBLE)
    if not found or size != 96:
        return None
    return bits[start:start + 96]