#!/usr/bin/env python3
"""Check captured KeeLoq presses (Richmond gate remote, HCS300) against the public key list.

Reads an ESPHome log from the CC1101 capture rig (keeloq-capture.yaml): the rig's own
"KEELOQ serial=... button=... hop=..." lines, and "Received Raw:" timings, which this script
decodes itself (ESPHome's decoder gives up if noise sits before the preamble).

For every key in keeloq_mfcodes_user it tries the learning schemes that work without a seed
(the HCS300 sends none): simple, normal, magic XOR 1, magic serial 1/2/3, each with the key
as written and byte-reversed. A key matches when two or more presses of one remote (any
buttons) all decrypt with the right button bits, the same discrimination bits, and a counter
that goes up by 1 to 32 from press to press. A single press whose discrimination bits equal
the serial's low 10 bits (Microchip's default) is only shown as a candidate.

Usage: python3 -I keeloq-check.py capture.log [keyfile]
       python3 -I keeloq-check.py --self-test
"""
import re
import statistics
import sys
from pathlib import Path

NLF = 0x3A5C742E
M32 = 0xFFFFFFFF
M64 = (1 << 64) - 1
KEYFILE = Path(__file__).with_name("keeloq_mfcodes_user")
TEST_KEYS = [(0x0123456789ABCDEF, 1, "made-up test key")]  # self-tests without the key list


def bit(x, n):
    return (x >> n) & 1


def encrypt(data, key):
    x = data
    for r in range(528):
        nlf = bit(NLF, bit(x, 1) | bit(x, 9) << 1 | bit(x, 20) << 2 | bit(x, 26) << 3 | bit(x, 31) << 4)
        x = (x >> 1) | ((bit(x, 0) ^ bit(x, 16) ^ bit(key, r & 63) ^ nlf) << 31)
    return x


def decrypt(data, key):
    x = data
    for r in range(528):
        nlf = bit(NLF, bit(x, 0) | bit(x, 8) << 1 | bit(x, 19) << 2 | bit(x, 25) << 3 | bit(x, 30) << 4)
        x = ((x << 1) & M32) | (bit(x, 31) ^ bit(x, 15) ^ bit(key, (15 - r) & 63) ^ nlf)
    return x


# Device key from the maker's key, as Flipper's keeloq_common.c does it. fix = button << 28 | serial
def simple(fix, mk):
    return mk


def normal(fix, mk):
    s = fix & 0x0FFFFFFF
    return decrypt(s | 0x60000000, mk) << 32 | decrypt(s | 0x20000000, mk)


def magic_xor_1(fix, mk):
    s = fix & 0x0FFFFFFF
    return (s << 32 | s) ^ mk


def magic_serial_1(fix, mk):
    return (mk & M32) | ((fix << 40) & M64) | ((((fix & 0xFF) + ((fix >> 8) & 0xFF)) & 0xFF) << 32)


def magic_serial_2(fix, mk):
    return int.from_bytes(fix.to_bytes(4, "little"), "big") << 32 | (mk & M32)


def magic_serial_3(fix, mk):
    return (mk & 0xFFFFFFFFFF000000) | (fix & 0xFFFFFF)


SCHEMES = [simple, normal, magic_xor_1, magic_serial_1, magic_serial_2, magic_serial_3]


def load_keys(path):
    keys = []
    for line in Path(path).read_text().splitlines():
        m = re.match(r"([0-9A-Fa-f]{16}):(\d+):(.+)", line.strip())
        if m:
            keys.append((int(m[1], 16), int(m[2]), m[3]))
    return keys


def decode_raw(durations):
    """Find KeeLoq frames in raw timings (+mark, -space, µs): 12-pulse preamble, 10 TE gap, 66 PWM bits.
    11 preamble pulses are enough: noise just before the frame can merge with the first one."""
    frames = []
    i = 0
    while i < len(durations) - 22:
        pre = durations[i : i + 21]
        te = statistics.median(abs(d) for d in pre)
        if not (60 < te < 800 and pre[0] > 0 and all(0.5 * te < abs(d) < 1.6 * te for d in pre)):
            i += 1
            continue
        j = i + 21  # the long header space; a longer preamble matches further along
        if j >= len(durations) or durations[j] > -6 * te:
            i += 1
            continue
        bits = []
        k = j + 1
        while len(bits) < 66 and k + 1 <= len(durations):
            mark = durations[k]
            space = -durations[k + 1] if k + 1 < len(durations) else 10 * te
            if mark <= 0 or not (0.5 * te < mark < 2.6 * te):
                break
            bits.append(1 if mark < 1.5 * te else 0)  # short mark = 1, long mark = 0
            k += 2
            if space > 4 * te:
                break
        if len(bits) == 66:
            v = sum(b << n for n, b in enumerate(bits))
            fix = (v >> 32) & M32
            frames.append({"serial": fix & 0x0FFFFFFF, "button": fix >> 28, "hop": v & M32, "te": round(te)})
            i = k
        else:
            i += 1
    return frames


def read_log(path):
    frames, raw = [], None
    for line in Path(path).read_text(errors="replace").splitlines():
        m = re.search(r"KEELOQ serial=(?:0x)?([0-9a-fA-F]+) button=(?:0x)?([0-9a-fA-F]+) hop=(?:0x)?([0-9a-fA-F]+)", line)
        if m:
            frames.append({"serial": int(m[1], 16), "button": int(m[2], 16), "hop": int(m[3], 16), "te": None})
            continue
        if "remote.raw" not in line:
            continue
        text = line.rsplit("]: ", 1)[-1]
        if "Received Raw:" in text:
            if raw:
                frames += decode_raw(raw)
            raw = []
            text = text.split("Received Raw:", 1)[1]
        if raw is not None:
            raw += [int(n) for n in re.findall(r"-?\d+", text)]
    if raw:
        frames += decode_raw(raw)
    # A held button repeats the same frame: keep each one once, in order
    seen, unique = set(), []
    for f in frames:
        k = (f["serial"], f["button"], f["hop"])
        if k not in seen:
            seen.add(k)
            unique.append(f)
    return unique


def check(frames, keys):
    """Return (confirmed, candidates). One press alone is only a candidate: with about 1,500
    key and scheme pairs, a random press passes the 14-bit check about 1 time in 11."""
    confirmed, candidates = [], []
    for mk, ktype, name in keys:
        for order, key in (("as written", mk), ("byte-reversed", int.from_bytes(mk.to_bytes(8, "big"), "little"))):
            for scheme in SCHEMES:
                by_serial = {}  # the HCS300 has one counter for all its buttons
                for f in frames:
                    fix = f["button"] << 28 | f["serial"]
                    d = decrypt(f["hop"], scheme(fix, key))
                    by_serial.setdefault(f["serial"], []).append((d >> 28 == f["button"], (d >> 16) & 0x3FF, d & 0xFFFF))
                for serial, presses in by_serial.items():
                    hit = (name, ktype, scheme.__name__, order, serial, [p[2] for p in presses])
                    counters = hit[5]
                    if (
                        len(presses) >= 2
                        and all(p[0] for p in presses)
                        and len({p[1] for p in presses}) == 1
                        and all(0 < (b - a) % 0x10000 <= 32 for a, b in zip(counters, counters[1:]))
                    ):
                        confirmed.append(hit)
                    elif any(p[0] and p[1] == serial & 0x3FF for p in presses):
                        candidates.append(hit)
    return confirmed, candidates


def report(frames, keys):
    if not frames:
        print("🔴 no KeeLoq frames found in the log")
        return False
    print(f"🌱 {len(frames)} distinct press(es), {len(keys)} keys × {len(SCHEMES)} schemes × 2 byte orders")
    for f in frames:
        te = f" TE {f['te']} µs" if f["te"] else ""
        print(f"🪵 serial 0x{f['serial']:07x} button 0x{f['button']:x} hop 0x{f['hop']:08x}{te}")
    confirmed, candidates = check(frames, keys)
    for name, ktype, scheme, order, serial, counters in confirmed:
        print(f"🟢 MATCH {name} (type {ktype}), {scheme}, key {order}: serial 0x{serial:07x}, counters {counters}")
    for name, ktype, scheme, order, serial, counters in candidates:
        print(f"🟠 one press fits {name} (type {ktype}), {scheme}, key {order}: serial 0x{serial:07x}; chance unless a second press fits too")
    if not confirmed:
        print("🔴 no key in the list decrypts two presses of one remote")
    return bool(confirmed)


def self_test():
    # Published KeeLoq test vector
    assert encrypt(0xF741E2DB, 0x5CEC6701B79FD949) == 0xE44F4CDF, "encrypt vector"
    assert decrypt(0xE44F4CDF, 0x5CEC6701B79FD949) == 0xF741E2DB, "decrypt vector"
    print("🟢 KeeLoq test vector")
    keys = load_keys(KEYFILE) if KEYFILE.exists() else TEST_KEYS
    mk, _, name = keys[len(keys) // 2]
    serial, button = 0x0ABCDEF, 0x2
    dev = normal(button << 28 | serial, mk)
    frames = []
    for cnt in (0x0101, 0x0102):
        hop = encrypt(button << 28 | (serial & 0x3FF) << 16 | cnt, dev)
        v = (button << 28 | serial) << 32 | hop
        raw = [400, -400] * 11 + [400, -4000]
        for n in range(66):
            raw += [400, -800] if bit(v, n) else [800, -400]
        raw[-1] = -15000
        frames += decode_raw([120, -90, 300] + raw)  # noise before the preamble
    assert len(frames) == 2 and frames[0]["serial"] == serial, f"raw decode: {frames}"
    print("🟢 raw decode of two synthetic presses")
    confirmed, _ = check(frames, keys)
    assert [(h[0], h[2], h[3]) for h in confirmed] == [(name, "normal", "as written")], f"wrong matches: {confirmed}"
    print(f"🟢 two synthetic presses made with '{name}' (normal learning) match it, and nothing else")
    random_presses = [{"serial": 0x0ABCDEF, "button": 2, "hop": encrypt(n, 0x1234567890ABCDEF), "te": None} for n in range(4)]
    assert not check(random_presses, keys)[0], "random presses matched"
    print("🟢 4 presses from an unknown key match nothing")
    print("\n🟢 SUCCESS")


if __name__ == "__main__":
    if sys.argv[1:2] == ["--self-test"]:
        self_test()
        sys.exit(0)
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(2)
    ok = report(read_log(sys.argv[1]), load_keys(sys.argv[2] if len(sys.argv) > 2 else KEYFILE))
    print("\n🟢 SUCCESS: a known key fits" if ok else "\n🔴 FAILED: no known key fits")
    sys.exit(0 if ok else 1)