pm3py/tests/test_sim_advanced.py

"""Tests for advanced applications: fuzzer, relay, replay."""
import asyncio
import pytest

from pm3py.sim.frame import RFFrame
from pm3py.sim.medium import SoftwareMedium
from pm3py.sim.iso14443a import Tag14443A_3, Reader14443A
from pm3py.sim.fuzzer import MutationFuzzer, GrammarFuzzer
from pm3py.sim.relay import RelayTransponder, MitMProxy
from pm3py.sim.replay import TraceRecorder, TraceReplayer, TraceEntry


def run(coro):
    return asyncio.get_event_loop().run_until_complete(coro)


# ---------------------------------------------------------------------------
# MutationFuzzer
# ---------------------------------------------------------------------------

class TestMutationFuzzer:
    def test_mutate_produces_different_frame(self):
        seed = RFFrame.from_hex("26")
        fuzzer = MutationFuzzer([seed], seed=42)
        mutated = fuzzer.mutate(seed)
        # At least sometimes should differ (with enough attempts)
        found_different = False
        for _ in range(20):
            m = fuzzer.mutate(seed)
            if m.data != seed.data or m.bit_count != seed.bit_count:
                found_different = True
                break
        assert found_different

    def test_generate_yields_frames(self):
        seeds = [RFFrame.from_hex("26"), RFFrame.from_hex("9320")]
        fuzzer = MutationFuzzer(seeds, seed=42)
        frames = list(fuzzer.generate(count=10))
        assert len(frames) == 10
        assert all(isinstance(f, RFFrame) for f in frames)

    def test_deterministic_with_seed(self):
        seeds = [RFFrame.from_hex("A5B6")]
        f1 = list(MutationFuzzer(seeds, seed=123).generate(5))
        f2 = list(MutationFuzzer(seeds, seed=123).generate(5))
        assert [f.data for f in f1] == [f.data for f in f2]

    def test_strategies_include_flip_truncate_extend(self):
        fuzzer = MutationFuzzer([RFFrame.from_hex("DEADBEEF")], seed=0)
        strategies_seen = set()
        for _ in range(100):
            frame = RFFrame.from_hex("DEADBEEF")
            m = fuzzer.mutate(frame)
            if m.bit_count < 32:
                strategies_seen.add("truncate")
            elif m.bit_count > 32:
                strategies_seen.add("extend")
            elif m.data != b"\xDE\xAD\xBE\xEF":
                strategies_seen.add("flip")
        assert len(strategies_seen) >= 2  # at least 2 strategies observed


# ---------------------------------------------------------------------------
# GrammarFuzzer
# ---------------------------------------------------------------------------

class TestGrammarFuzzer:
    def test_generate_14443a_reqa(self):
        fuzzer = GrammarFuzzer("14443a")
        frame = fuzzer.generate("REQA")
        assert frame is not None
        assert frame.bit_count > 0

    def test_generate_with_overrides(self):
        fuzzer = GrammarFuzzer("14443a")
        frame = fuzzer.generate("SELECT", uid=b"\xFF\xFF\xFF\xFF")
        assert frame is not None

    def test_unknown_command_returns_none(self):
        fuzzer = GrammarFuzzer("14443a")
        frame = fuzzer.generate("NONEXISTENT")
        assert frame is None


# ---------------------------------------------------------------------------
# RelayTransponder
# ---------------------------------------------------------------------------

class TestRelayTransponder:
    def test_relay_forwards_frames(self):
        """RelayTransponder forwards reader commands to upstream reader."""
        # Set up "real card" on upstream medium
        upstream_medium = SoftwareMedium()
        uid = b"\x01\x02\x03\x04"
        run(upstream_medium.attach(Tag14443A_3(uid=uid, sak=0x08)))
        upstream_reader = Reader14443A(upstream_medium)

        # Set up relay on local medium
        relay = RelayTransponder(upstream_reader)
        local_medium = SoftwareMedium()
        run(local_medium.attach(relay))

        # Send REQA through local medium
        run(local_medium.transmit_reader(RFFrame.from_hex("26")))
        resp = run(local_medium.receive_reader())
        # Should get ATQA back from the real card
        assert resp is not None
        assert len(resp.data) == 2  # ATQA


# ---------------------------------------------------------------------------
# MitMProxy
# ---------------------------------------------------------------------------

class TestMitMProxy:
    def test_proxy_logs_traffic(self):
        reader_medium = SoftwareMedium()
        tag_medium = SoftwareMedium()
        run(tag_medium.attach(Tag14443A_3(uid=b"\x01\x02\x03\x04")))

        proxy = MitMProxy(reader_medium, tag_medium)

        # Simulate reader → proxy → tag
        run(proxy.forward_to_tag(RFFrame.from_hex("26")))
        assert len(proxy.trace) == 2  # request + response

    def test_proxy_with_intercept(self):
        reader_medium = SoftwareMedium()
        tag_medium = SoftwareMedium()
        run(tag_medium.attach(Tag14443A_3(uid=b"\x01\x02\x03\x04")))

        modified = []
        def intercept(frame, direction):
            modified.append(direction)
            return frame  # pass through

        proxy = MitMProxy(reader_medium, tag_medium, intercept=intercept)
        run(proxy.forward_to_tag(RFFrame.from_hex("26")))
        assert "reader" in modified
        assert "tag" in modified


# ---------------------------------------------------------------------------
# TraceRecorder / TraceReplayer
# ---------------------------------------------------------------------------

class TestTraceRecorder:
    def test_records_traffic(self):
        medium = SoftwareMedium()
        run(medium.attach(Tag14443A_3(uid=b"\x01\x02\x03\x04")))
        recorder = TraceRecorder(medium)

        # Exchange through recorder
        run(recorder.transmit_reader(RFFrame.from_hex("26")))
        run(recorder.receive_reader())

        assert len(recorder.trace) >= 1

    def test_trace_entries_have_direction(self):
        medium = SoftwareMedium()
        run(medium.attach(Tag14443A_3(uid=b"\x01\x02\x03\x04")))
        recorder = TraceRecorder(medium)

        run(recorder.transmit_reader(RFFrame.from_hex("26")))
        resp = run(recorder.receive_reader())

        directions = [e.direction for e in recorder.trace]
        assert "reader" in directions
        if resp is not None:
            assert "tag" in directions


class TestTraceReplayer:
    def test_replay_against_tag(self):
        """Record a trace, then replay it against a fresh tag."""
        # Record
        medium1 = SoftwareMedium()
        run(medium1.attach(Tag14443A_3(uid=b"\x01\x02\x03\x04")))
        recorder = TraceRecorder(medium1)
        run(recorder.transmit_reader(RFFrame.from_hex("26")))
        run(recorder.receive_reader())

        # Replay against fresh tag
        medium2 = SoftwareMedium()
        run(medium2.attach(Tag14443A_3(uid=b"\x01\x02\x03\x04")))
        replayer = TraceReplayer(recorder.trace, medium2)
        results = run(replayer.replay())
        assert len(results) > 0
        assert all(r["matched"] for r in results)