Adds a third export alongside pull history and achievements: the cartridges, modules, character levels and equipped console boards the game sends at login. Shaped as added files plus a handful of one-line calls, so this repo can keep merging upstream: - decoder/gear.py: self-contained bit-level decoder over a list of datagrams. - export/gear_export.py: the nte-gear-export envelope. - live_capture/gear_collector.py: UDP flow accumulation, the mid-capture decode, the fallback scan, the console line and the export file. Gear is reported during the capture rather than only at the stop. The burst arrives in a rush and then stops, so "this flow went quiet" is the only mid-stream marker available; a flow is retried only once it has grown, since decoding per packet would be quadratic in flow size. Flows are tried largest first and the first one that decodes wins - picking the single biggest was wrong, history traffic can outweigh a short login burst. Verified against a live capture: 817 items (310 cartridges, 507 modules), 20 characters, 13 loadouts, every bucket matching the in-game counts exactly. 26 synthetic tests, no real capture used as a fixture; 149 total. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
747 lines
27 KiB
Python
747 lines
27 KiB
Python
"""Gear decoding, driven entirely by synthetic packets.
|
|
|
|
The wire format is fully understood, so every packet here is constructed from
|
|
the spec rather than captured: a 120-bit header, payload bits, two 1 bits and
|
|
zero padding. No real capture is used - see tests/fixtures/README.md.
|
|
"""
|
|
|
|
import contextlib
|
|
import io
|
|
import json
|
|
import os
|
|
import struct
|
|
import tempfile
|
|
import unittest
|
|
|
|
from nte_history_exporter import console
|
|
|
|
from nte_history_exporter import __version__
|
|
from nte_history_exporter.decoder.gear import (
|
|
HEADER_BITS,
|
|
extract_gear,
|
|
payload_end_bit,
|
|
splice_datagrams,
|
|
)
|
|
from nte_history_exporter.export.gear_export import build_gear_export_json
|
|
from nte_history_exporter.live_capture import gear_collector, runner
|
|
from nte_history_exporter.live_capture.windows_raw import ParsedIpPacket
|
|
|
|
from tests import support
|
|
|
|
|
|
# --- synthetic wire construction -------------------------------------------
|
|
|
|
|
|
def string_frame(text: bytes) -> bytes:
|
|
"""A self-validating <u32 len><ascii\\0> frame."""
|
|
return struct.pack("<I", len(text) + 1) + text + b"\0"
|
|
|
|
|
|
def stat_entry(name: bytes, value: float) -> bytes:
|
|
"""<u32 len><name\\0><4 zero bytes><f32 value>."""
|
|
return string_frame(name) + b"\0\0\0\0" + struct.pack("<f", value)
|
|
|
|
|
|
def poke(buf: bytearray, byte_offset: int, raw: bytes) -> None:
|
|
if len(buf) < byte_offset + len(raw):
|
|
buf.extend(b"\0" * (byte_offset + len(raw) - len(buf)))
|
|
buf[byte_offset : byte_offset + len(raw)] = raw
|
|
|
|
|
|
def poke_u32_at_bit(buf: bytearray, bitpos: int, value: int) -> None:
|
|
"""Write a little-endian u32 at an arbitrary bit offset."""
|
|
needed = (bitpos + 32 + 7) // 8 + 1
|
|
if len(buf) < needed:
|
|
buf.extend(b"\0" * (needed - len(buf)))
|
|
whole = int.from_bytes(buf, "little")
|
|
whole &= ~(0xFFFFFFFF << bitpos)
|
|
whole |= value << bitpos
|
|
packed = whole.to_bytes(len(buf) + 8, "little")
|
|
buf[:] = packed[: len(buf)]
|
|
|
|
|
|
def item_record(
|
|
item_id: bytes,
|
|
*,
|
|
level: int = 20,
|
|
entries=(),
|
|
instance: bytes = b"\x11\x22\x33\x44\x55\x66\x77\x88",
|
|
owner: bytes | None = None,
|
|
) -> bytes:
|
|
"""One item record laid out at the documented offsets past the id frame."""
|
|
body = bytearray(b"\0" * 64)
|
|
poke(body, 5, instance)
|
|
# Level is a u32 at +38 bytes; the owner reference starts at +39 bits-wise
|
|
# (312 bits) and its first three bytes are always zero, so the two fields
|
|
# coexist exactly as they do on the wire.
|
|
poke(body, 38, struct.pack("<I", level))
|
|
if owner is not None:
|
|
assert owner[:3] == b"\0\0\0", "owner reference starts with three zero bytes"
|
|
poke(body, 39, owner)
|
|
chain = b"".join(stat_entry(name, value) for name, value in entries)
|
|
return string_frame(item_id) + bytes(body) + chain
|
|
|
|
|
|
def character_block(name: bytes, *, level: int, breakthroughs: int) -> bytes:
|
|
"""An ability frame with the progress fields at their negative offsets."""
|
|
lead = 128 # bytes of room for the fields that sit before the frame
|
|
buf = bytearray(b"\0" * lead)
|
|
frame = string_frame(b"GA_" + name + b"_Melee")
|
|
block_bit = lead * 8
|
|
poke_u32_at_bit(buf, block_bit - 466, 67806)
|
|
poke_u32_at_bit(buf, block_bit - 434, level)
|
|
poke_u32_at_bit(buf, block_bit - 402, breakthroughs)
|
|
return bytes(buf) + frame
|
|
|
|
|
|
def loadout_cell(item_id: bytes, row: int, column: int) -> bytes:
|
|
"""An id frame quoted with board coordinates and no stat chain."""
|
|
body = bytearray(b"\0" * 32)
|
|
# Coordinates sit 96 and 128 bits past the id and are stored doubled.
|
|
poke(body, 12, struct.pack("<I", row * 2))
|
|
poke(body, 16, struct.pack("<I", column * 2))
|
|
return string_frame(item_id) + bytes(body)
|
|
|
|
|
|
def pack_datagrams(payload: bytes, *, chunk_bits: int = 1 << 20, first_seq: int = 1):
|
|
"""Split a payload bit stream into framed datagrams."""
|
|
total_bits = len(payload) * 8
|
|
whole = int.from_bytes(payload, "little")
|
|
datagrams, start, seq = [], 0, first_seq
|
|
while start < total_bits:
|
|
nbits = min(chunk_bits, total_bits - start)
|
|
chunk = (whole >> start) & ((1 << nbits) - 1)
|
|
|
|
header = bytearray(15)
|
|
header[0:3] = b"\xaa\xbb\xcc"
|
|
# 14-bit packet sequence at bit 24.
|
|
header[3] = seq & 0xFF
|
|
header[4] = (seq >> 8) & 0x3F
|
|
header[12:15] = b"\x64\x06\xee"
|
|
|
|
# payload bits, then the two 1 bits Unreal terminates with.
|
|
value = int.from_bytes(header, "little")
|
|
value |= chunk << HEADER_BITS
|
|
value |= 0b11 << (HEADER_BITS + nbits)
|
|
width = (HEADER_BITS + nbits + 2 + 7) // 8
|
|
datagrams.append(value.to_bytes(width, "little"))
|
|
|
|
start += nbits
|
|
seq = (seq + 1) & 0x3FFF
|
|
return datagrams
|
|
|
|
|
|
MODULE_ENTRIES = (
|
|
(b"HPMaxAdd", 0.0),
|
|
(b"AtkAdd", 0.0),
|
|
(b"CritDamageBase", 0.06),
|
|
(b"CritBase", 0.03),
|
|
(b"AtkUp", 0.04),
|
|
(b"HPMaxUp", 0.04),
|
|
)
|
|
|
|
CARTRIDGE_ENTRIES = (
|
|
(b"AtkUp", 0.0),
|
|
(b"CritDamageBase", 0.06),
|
|
(b"CritBase", 0.03),
|
|
(b"DefUp", 0.07),
|
|
(b"HPMaxUp", 0.04),
|
|
)
|
|
|
|
OWNER = b"\0\0\0\x0d\x0d\x20\x3d\x6d\x00\x00\x00\x00"
|
|
|
|
|
|
class FramingTests(unittest.TestCase):
|
|
def test_payload_end_bit_finds_the_double_terminator(self):
|
|
[datagram] = pack_datagrams(b"\x01\x02\x03\x04")
|
|
self.assertEqual(payload_end_bit(datagram), HEADER_BITS + 32)
|
|
|
|
def test_splice_round_trips_a_payload_across_datagrams(self):
|
|
payload = bytes(range(256)) * 3
|
|
# A chunk width that is not a multiple of 8 forces the accumulator to
|
|
# rejoin the stream at a bit offset, as the real capture does.
|
|
datagrams = pack_datagrams(payload, chunk_bits=333)
|
|
self.assertGreater(len(datagrams), 1)
|
|
self.assertEqual(splice_datagrams(datagrams)[: len(payload)], payload)
|
|
|
|
def test_pure_acks_are_skipped(self):
|
|
payload = b"\x01\x02\x03\x04\x05\x06\x07\x08"
|
|
ack = bytes(12)
|
|
datagrams = pack_datagrams(payload)
|
|
self.assertEqual(splice_datagrams([ack] + datagrams + [ack])[:8], payload)
|
|
|
|
|
|
class ItemRecordTests(unittest.TestCase):
|
|
def decode(self, payload: bytes, **kwargs):
|
|
return extract_gear(pack_datagrams(b"\0" * 8 + payload, **kwargs))
|
|
|
|
def test_module_record_decodes(self):
|
|
snapshot = self.decode(
|
|
item_record(b"cell3_style6_1_Orange", level=20, entries=MODULE_ENTRIES)
|
|
)
|
|
[item] = snapshot.items
|
|
self.assertEqual(item.kind, "module")
|
|
self.assertEqual(item.item_id, "cell3_style6_1_Orange")
|
|
self.assertEqual(item.shape, "cell3_style6")
|
|
self.assertEqual(item.module_type, "III")
|
|
self.assertEqual(item.level, 20)
|
|
self.assertEqual(item.rarity, "orange")
|
|
self.assertEqual(item.instance, "1122334455667788")
|
|
self.assertIsNone(item.owner_group)
|
|
# Both mains are computed from the cell count, never transmitted.
|
|
self.assertEqual(dict(item.main_stats), {"HPMaxAdd": 840.0, "AtkAdd": 63.0})
|
|
self.assertEqual(len(item.substats), 4)
|
|
|
|
def test_cartridge_record_decodes_with_its_set(self):
|
|
snapshot = self.decode(
|
|
item_record(b"Incantation_orange", level=20, entries=CARTRIDGE_ENTRIES)
|
|
)
|
|
[item] = snapshot.items
|
|
self.assertEqual(item.kind, "cartridge")
|
|
self.assertIsNotNone(item.set_name)
|
|
self.assertEqual(len(item.main_stats), 1)
|
|
self.assertEqual(len(item.substats), 4)
|
|
|
|
def test_owner_reference_is_read_when_equipped(self):
|
|
snapshot = self.decode(
|
|
item_record(
|
|
b"cell3_style6_1_Orange",
|
|
level=20,
|
|
entries=MODULE_ENTRIES,
|
|
owner=OWNER,
|
|
)
|
|
)
|
|
[item] = snapshot.items
|
|
self.assertEqual(item.owner_group, OWNER.hex())
|
|
# The owner field overlaps the level word without corrupting it.
|
|
self.assertEqual(item.level, 20)
|
|
|
|
def test_record_spanning_a_datagram_boundary_decodes(self):
|
|
record = item_record(
|
|
b"cell3_style6_1_Orange", level=20, entries=MODULE_ENTRIES, owner=OWNER
|
|
)
|
|
# Chunks far smaller than the record, at a width that is not a whole
|
|
# number of bytes: every field lands across at least one boundary.
|
|
snapshot = self.decode(record, chunk_bits=101)
|
|
[item] = snapshot.items
|
|
self.assertEqual(item.item_id, "cell3_style6_1_Orange")
|
|
self.assertEqual(item.level, 20)
|
|
self.assertEqual(item.owner_group, OWNER.hex())
|
|
self.assertEqual(dict(item.main_stats), {"HPMaxAdd": 840.0, "AtkAdd": 63.0})
|
|
|
|
def test_unframed_id_is_ignored(self):
|
|
# Same bytes, wrong length prefix: the frame check must reject it.
|
|
record = bytearray(
|
|
item_record(b"cell3_style6_1_Orange", entries=MODULE_ENTRIES)
|
|
)
|
|
record[0:4] = struct.pack("<I", 99)
|
|
self.assertEqual(self.decode(bytes(record)).items, ())
|
|
|
|
|
|
class ValidationTests(unittest.TestCase):
|
|
def decode_one(self, payload: bytes):
|
|
return extract_gear(pack_datagrams(b"\0" * 8 + payload)).items
|
|
|
|
def test_main_stat_after_a_substat_is_rejected(self):
|
|
scrambled = (
|
|
(b"HPMaxAdd", 0.0),
|
|
(b"CritDamageBase", 0.06),
|
|
(b"AtkAdd", 0.0),
|
|
(b"CritBase", 0.03),
|
|
(b"AtkUp", 0.04),
|
|
(b"HPMaxUp", 0.04),
|
|
)
|
|
self.assertEqual(
|
|
self.decode_one(
|
|
item_record(b"cell3_style6_1_Orange", entries=scrambled)
|
|
),
|
|
(),
|
|
)
|
|
|
|
def test_off_table_substat_value_is_rejected(self):
|
|
corrupt = MODULE_ENTRIES[:-1] + ((b"HPMaxUp", 0.0413),)
|
|
self.assertEqual(
|
|
self.decode_one(item_record(b"cell3_style6_1_Orange", entries=corrupt)),
|
|
(),
|
|
)
|
|
|
|
def test_short_entry_chain_is_rejected(self):
|
|
# Five entries is a cartridge count; a module needs six.
|
|
self.assertEqual(
|
|
self.decode_one(
|
|
item_record(b"cell3_style6_1_Orange", entries=MODULE_ENTRIES[:5])
|
|
),
|
|
(),
|
|
)
|
|
|
|
|
|
class CharacterTests(unittest.TestCase):
|
|
def test_level_and_breakthroughs_are_read(self):
|
|
payload = character_block(b"Zankou", level=70, breakthroughs=5)
|
|
[character] = extract_gear(pack_datagrams(payload)).characters
|
|
self.assertEqual(character.key, "Zankou")
|
|
self.assertEqual(character.level, 70)
|
|
self.assertEqual(character.breakthroughs, 5)
|
|
|
|
def test_progress_without_the_marker_is_not_reported(self):
|
|
payload = bytearray(character_block(b"Zankou", level=70, breakthroughs=5))
|
|
poke_u32_at_bit(payload, 128 * 8 - 466, 1) # marker no longer 67806
|
|
[character] = extract_gear(pack_datagrams(bytes(payload))).characters
|
|
self.assertIsNone(character.level)
|
|
self.assertIsNone(character.breakthroughs)
|
|
|
|
def test_repeated_block_keeps_the_first_valid_reading(self):
|
|
block = character_block(b"Zankou", level=70, breakthroughs=5)
|
|
blank = bytearray(character_block(b"Zankou", level=70, breakthroughs=5))
|
|
poke_u32_at_bit(blank, 128 * 8 - 466, 1)
|
|
payload = block + b"\0" * 64 + bytes(blank)
|
|
[character] = extract_gear(pack_datagrams(payload)).characters
|
|
self.assertEqual(character.level, 70)
|
|
self.assertEqual(character.breakthroughs, 5)
|
|
|
|
|
|
class LoadoutTests(unittest.TestCase):
|
|
def test_board_cells_follow_their_character_and_are_halved(self):
|
|
payload = (
|
|
character_block(b"Zankou", level=70, breakthroughs=5)
|
|
+ b"\0" * 64
|
|
+ loadout_cell(b"cell3_style6_1_Orange", 1, 2)
|
|
+ b"\0" * 64
|
|
+ loadout_cell(b"cell4_style2_1_Orange", 3, 4)
|
|
)
|
|
snapshot = extract_gear(pack_datagrams(payload))
|
|
[character] = snapshot.characters
|
|
# A transposed board renders correctly on a symmetric grid, so the
|
|
# assertion has to pin row and column separately.
|
|
self.assertEqual(character.board, ((1, 2), (3, 4)))
|
|
self.assertEqual(snapshot.characters_with_loadouts, 1)
|
|
# Quoted ids are references, not records.
|
|
self.assertEqual(snapshot.items, ())
|
|
|
|
|
|
class GearExportContractTests(unittest.TestCase):
|
|
def build(self):
|
|
payload = (
|
|
character_block(b"Zankou", level=70, breakthroughs=5)
|
|
+ b"\0" * 64
|
|
+ loadout_cell(b"cell3_style6_1_Orange", 1, 2)
|
|
+ b"\0" * 64
|
|
+ item_record(
|
|
b"cell3_style6_1_Orange",
|
|
level=20,
|
|
entries=MODULE_ENTRIES,
|
|
owner=OWNER,
|
|
)
|
|
+ b"\0" * 64
|
|
+ item_record(
|
|
b"Incantation_orange",
|
|
level=20,
|
|
entries=CARTRIDGE_ENTRIES,
|
|
instance=b"\x99\xaa\xbb\xcc\xdd\xee\xff\x01",
|
|
)
|
|
)
|
|
snapshot = extract_gear(pack_datagrams(payload))
|
|
return snapshot, build_gear_export_json(
|
|
snapshot,
|
|
source="live_capture",
|
|
capture_source="libpcap",
|
|
user_uid=" 100200300 ",
|
|
server_id="1",
|
|
)
|
|
|
|
def test_envelope(self):
|
|
_snapshot, export = self.build()
|
|
self.assertEqual(export["format"], "nte-gear-export")
|
|
self.assertEqual(export["format_version"], 1)
|
|
self.assertEqual(export["source"], "live_capture")
|
|
self.assertEqual(export["capture_source"], "libpcap")
|
|
self.assertEqual(export["exporter"]["version"], __version__)
|
|
self.assertEqual(export["user_uid"], "100200300")
|
|
self.assertEqual(export["server_id"], "1")
|
|
|
|
def test_scan_counts_match_the_snapshot(self):
|
|
snapshot, export = self.build()
|
|
self.assertEqual(
|
|
export["scan"],
|
|
{
|
|
"cartridges": 1,
|
|
"modules": 1,
|
|
"characters": 1,
|
|
"characters_with_loadouts": 1,
|
|
"warnings": [],
|
|
},
|
|
)
|
|
self.assertEqual(len(export["items"]), len(snapshot.items))
|
|
|
|
def test_item_records_carry_raw_stat_ids(self):
|
|
_snapshot, export = self.build()
|
|
module = next(i for i in export["items"] if i["kind"] == "module")
|
|
self.assertEqual(module["shape"], "cell3_style6")
|
|
self.assertEqual(module["module_type"], "III")
|
|
self.assertEqual(module["owner_group"], OWNER.hex())
|
|
self.assertIn(
|
|
{"stat": "HPMaxAdd", "value": 840.0}, module["main_stats"]
|
|
)
|
|
self.assertTrue(all("stat" in s for s in module["substats"]))
|
|
cartridge = next(i for i in export["items"] if i["kind"] == "cartridge")
|
|
self.assertIn("set", cartridge)
|
|
self.assertIsNone(cartridge["owner_group"])
|
|
|
|
def test_character_board_and_unresolved_owner_group(self):
|
|
_snapshot, export = self.build()
|
|
[character] = export["characters"]
|
|
self.assertEqual(character["key"], "Zankou")
|
|
self.assertEqual(character["level"], 70)
|
|
self.assertEqual(character["breakthroughs"], 5)
|
|
self.assertEqual(character["board"], [[1, 2]])
|
|
# The capture cannot join an owner group to a name; emitting a guess
|
|
# would be worse than emitting the truth.
|
|
self.assertIsNone(character["owner_group"])
|
|
|
|
|
|
# --- runner wiring ---------------------------------------------------------
|
|
|
|
|
|
class _StubCapture:
|
|
name = "libpcap"
|
|
detail = "documentation range"
|
|
fallback_reason = ""
|
|
|
|
def __init__(self, packets):
|
|
self._packets = packets
|
|
|
|
def packets(self):
|
|
return iter(self._packets)
|
|
|
|
def stats(self):
|
|
return None
|
|
|
|
def close(self):
|
|
return None
|
|
|
|
|
|
LOCAL_IP = "192.0.2.10"
|
|
REMOTE_IP = "198.51.100.20"
|
|
|
|
|
|
class RunnerWiringTests(unittest.TestCase):
|
|
"""The gear export must survive being the only thing captured."""
|
|
|
|
def run_capture(self, *, user_uid=None):
|
|
payload = (
|
|
character_block(b"Zankou", level=70, breakthroughs=5)
|
|
+ b"\0" * 64
|
|
+ loadout_cell(b"cell3_style6_1_Orange", 1, 2)
|
|
+ b"\0" * 64
|
|
+ item_record(
|
|
b"cell3_style6_1_Orange",
|
|
level=20,
|
|
entries=MODULE_ENTRIES,
|
|
owner=OWNER,
|
|
)
|
|
)
|
|
packets = [
|
|
ParsedIpPacket(
|
|
src_ip=REMOTE_IP,
|
|
dst_ip=LOCAL_IP,
|
|
src_port=30166,
|
|
dst_port=49610,
|
|
payload=datagram,
|
|
protocol="udp",
|
|
)
|
|
for datagram in pack_datagrams(payload, chunk_bits=277)
|
|
]
|
|
# An outbound datagram and a second, quieter inbound flow: the decoder
|
|
# must pick the busiest inbound flow and ignore the rest.
|
|
packets.append(
|
|
ParsedIpPacket(
|
|
src_ip=LOCAL_IP,
|
|
dst_ip=REMOTE_IP,
|
|
src_port=49610,
|
|
dst_port=30166,
|
|
payload=b"\x00" * 32,
|
|
protocol="udp",
|
|
)
|
|
)
|
|
packets.append(
|
|
ParsedIpPacket(
|
|
src_ip=REMOTE_IP,
|
|
dst_ip=LOCAL_IP,
|
|
src_port=30167,
|
|
dst_port=49611,
|
|
payload=b"\x00" * 16,
|
|
protocol="udp",
|
|
)
|
|
)
|
|
|
|
prompted = []
|
|
|
|
def fake_prompt_user_uid():
|
|
prompted.append(True)
|
|
return "100200300"
|
|
|
|
with tempfile.TemporaryDirectory() as workdir:
|
|
previous = os.getcwd()
|
|
os.chdir(workdir)
|
|
patched_backend = runner.open_capture_backend
|
|
patched_prompt = console.prompt_user_uid
|
|
runner.open_capture_backend = lambda ip, backend: _StubCapture(packets)
|
|
console.prompt_user_uid = fake_prompt_user_uid
|
|
try:
|
|
with contextlib.redirect_stdout(io.StringIO()) as out:
|
|
result = runner.run_live_capture(
|
|
interface_ip=LOCAL_IP, user_uid=user_uid
|
|
)
|
|
gear_path = result["gear_path"]
|
|
self.assertIsNotNone(gear_path)
|
|
export = json.loads(gear_path.read_text(encoding="utf-8"))
|
|
printed = out.getvalue()
|
|
finally:
|
|
runner.open_capture_backend = patched_backend
|
|
console.prompt_user_uid = patched_prompt
|
|
os.chdir(previous)
|
|
return result, export, printed, prompted
|
|
|
|
def test_gear_only_capture_still_writes_and_reports_its_export(self):
|
|
result, export, printed, prompted = self.run_capture()
|
|
self.assertEqual(result["exports"], [])
|
|
self.assertEqual(export["format"], "nte-gear-export")
|
|
self.assertEqual(export["scan"]["modules"], 1)
|
|
self.assertEqual(export["scan"]["characters_with_loadouts"], 1)
|
|
self.assertEqual(export["source"], "live_capture")
|
|
self.assertEqual(export["capture_source"], "libpcap")
|
|
# The summary branch for "no pull history" has to mention it, or the
|
|
# file is written and never named.
|
|
self.assertIn("Gear export complete.", printed)
|
|
self.assertIn(result["gear_path"].name, printed)
|
|
self.assertIn("Gear captured", printed)
|
|
# Gear alone must be enough to opt into the UID prompt.
|
|
self.assertEqual(prompted, [True])
|
|
self.assertEqual(export["user_uid"], "100200300")
|
|
self.assertTrue(result["gear_path"].name.startswith("100200300_Gear_"))
|
|
|
|
def test_supplied_uid_skips_the_prompt(self):
|
|
_result, export, _printed, prompted = self.run_capture(user_uid="999")
|
|
self.assertEqual(prompted, [])
|
|
self.assertEqual(export["user_uid"], "999")
|
|
|
|
|
|
def test_gear_export_is_named_alongside_pull_history_exports(self):
|
|
"""The second summary branch: gear must not vanish behind a banner."""
|
|
fixture = support.load_network_fixture()
|
|
local_ip = fixture["local_ip"]
|
|
payload = item_record(
|
|
b"cell3_style6_1_Orange", level=20, entries=MODULE_ENTRIES, owner=OWNER
|
|
)
|
|
packets = [
|
|
ParsedIpPacket(
|
|
src_ip=packet.src_ip,
|
|
dst_ip=packet.dst_ip,
|
|
src_port=packet.src_port,
|
|
dst_port=packet.dst_port,
|
|
payload=packet.payload,
|
|
protocol=packet.protocol,
|
|
)
|
|
for packet in support.fixture_packets("replay")
|
|
]
|
|
packets += [
|
|
ParsedIpPacket(
|
|
src_ip=REMOTE_IP,
|
|
dst_ip=local_ip,
|
|
src_port=30166,
|
|
dst_port=49610,
|
|
payload=datagram,
|
|
protocol="udp",
|
|
)
|
|
for datagram in pack_datagrams(payload, chunk_bits=277)
|
|
]
|
|
|
|
with tempfile.TemporaryDirectory() as workdir:
|
|
previous = os.getcwd()
|
|
os.chdir(workdir)
|
|
patched_backend = runner.open_capture_backend
|
|
patched_server = console.prompt_server_id
|
|
runner.open_capture_backend = lambda ip, backend: _StubCapture(packets)
|
|
console.prompt_server_id = lambda: ""
|
|
try:
|
|
with contextlib.redirect_stdout(io.StringIO()) as out:
|
|
result = runner.run_live_capture(
|
|
interface_ip=local_ip, user_uid="100200300"
|
|
)
|
|
printed = out.getvalue()
|
|
gear_path = result["gear_path"]
|
|
self.assertIsNotNone(gear_path)
|
|
export = json.loads(gear_path.read_text(encoding="utf-8"))
|
|
finally:
|
|
runner.open_capture_backend = patched_backend
|
|
console.prompt_server_id = patched_server
|
|
os.chdir(previous)
|
|
|
|
self.assertTrue(result["exports"], "the fixture should still yield a banner")
|
|
self.assertEqual(export["scan"]["modules"], 1)
|
|
self.assertIn(gear_path.name, printed)
|
|
|
|
def test_udp_flow_accumulation_is_bounded(self):
|
|
self.assertEqual(gear_collector.MAX_UDP_FLOW_BYTES, 8 * 1024 * 1024)
|
|
|
|
|
|
class _FakeClock:
|
|
"""Stands in for the `time` module so idleness is exact, not slept for."""
|
|
|
|
def __init__(self, now=1000.0):
|
|
self.now = now
|
|
|
|
def monotonic(self):
|
|
return self.now
|
|
|
|
def time(self):
|
|
return self.now
|
|
|
|
|
|
class _TickingCapture(_StubCapture):
|
|
"""Yields a burst, then the read-timeout ticks pcap emits when idle.
|
|
|
|
`pending` is how many items are still unyielded, which is what proves a
|
|
report happened during the capture rather than after it.
|
|
"""
|
|
|
|
def __init__(self, packets, clock, *, ticks=3, tick_seconds=5.0):
|
|
super().__init__(packets)
|
|
self.clock = clock
|
|
self.ticks = ticks
|
|
self.tick_seconds = tick_seconds
|
|
self.pending = 0
|
|
|
|
def packets(self):
|
|
queued = list(self._packets) + [None] * self.ticks
|
|
self.pending = len(queued)
|
|
for item in queued:
|
|
self.pending -= 1
|
|
if item is None:
|
|
self.clock.now += self.tick_seconds
|
|
yield item
|
|
|
|
|
|
class LiveGearReportingTests(unittest.TestCase):
|
|
"""Gear is reported while the capture runs, not only once it stops."""
|
|
|
|
def burst_packets(self):
|
|
payload = (
|
|
character_block(b"Zankou", level=70, breakthroughs=5)
|
|
+ b"\0" * 64
|
|
+ loadout_cell(b"cell3_style6_1_Orange", 1, 2)
|
|
+ b"\0" * 64
|
|
+ item_record(
|
|
b"cell3_style6_1_Orange",
|
|
level=20,
|
|
entries=MODULE_ENTRIES,
|
|
owner=OWNER,
|
|
)
|
|
)
|
|
return [
|
|
ParsedIpPacket(
|
|
src_ip=REMOTE_IP,
|
|
dst_ip=LOCAL_IP,
|
|
src_port=30166,
|
|
dst_port=49610,
|
|
payload=datagram,
|
|
protocol="udp",
|
|
)
|
|
for datagram in pack_datagrams(payload, chunk_bits=277)
|
|
]
|
|
|
|
def run_capture(self, capture_factory, *, min_flow_bytes=1):
|
|
"""Run one live capture against a stub, with the size gate relaxed.
|
|
|
|
Synthetic bursts are a few hundred bytes; the real gate exists to keep
|
|
chatter out, not to describe a minimum decodable size.
|
|
"""
|
|
clock = _FakeClock()
|
|
capture = capture_factory(clock)
|
|
decodes = []
|
|
real_extract = gear_collector.extract_gear
|
|
reported_pending = []
|
|
real_report = gear_collector.print_gear_captured
|
|
|
|
def counting_extract(grams):
|
|
decodes.append(len(grams))
|
|
return real_extract(grams)
|
|
|
|
def recording_report(*args, **kwargs):
|
|
reported_pending.append(capture.pending)
|
|
return real_report(*args, **kwargs)
|
|
|
|
with tempfile.TemporaryDirectory() as workdir:
|
|
previous = os.getcwd()
|
|
os.chdir(workdir)
|
|
patched_backend = runner.open_capture_backend
|
|
patched_time = gear_collector.time
|
|
runner.open_capture_backend = lambda ip, backend: capture
|
|
gear_collector.time = clock
|
|
gear_collector.extract_gear = counting_extract
|
|
gear_collector.print_gear_captured = recording_report
|
|
previous_min = gear_collector.GEAR_LIVE_MIN_FLOW_BYTES
|
|
gear_collector.GEAR_LIVE_MIN_FLOW_BYTES = min_flow_bytes
|
|
try:
|
|
with contextlib.redirect_stdout(io.StringIO()) as out:
|
|
result = runner.run_live_capture(
|
|
interface_ip=LOCAL_IP, user_uid="100200300"
|
|
)
|
|
printed = out.getvalue()
|
|
finally:
|
|
runner.open_capture_backend = patched_backend
|
|
gear_collector.time = patched_time
|
|
gear_collector.extract_gear = real_extract
|
|
gear_collector.print_gear_captured = real_report
|
|
gear_collector.GEAR_LIVE_MIN_FLOW_BYTES = previous_min
|
|
os.chdir(previous)
|
|
return result, printed, decodes, reported_pending
|
|
|
|
def test_gear_is_reported_before_the_capture_ends(self):
|
|
result, printed, _, reported_pending = self.run_capture(
|
|
lambda clock: _TickingCapture(self.burst_packets(), clock)
|
|
)
|
|
|
|
self.assertIn("Gear captured", printed)
|
|
self.assertEqual(len(reported_pending), 1, "reported exactly once")
|
|
self.assertGreater(
|
|
reported_pending[0],
|
|
0,
|
|
"the report must land while packets are still to come",
|
|
)
|
|
self.assertIsNotNone(result["gear_path"])
|
|
|
|
def test_a_flow_that_never_goes_idle_still_decodes_at_the_stop(self):
|
|
result, printed, _, reported_pending = self.run_capture(
|
|
lambda clock: _TickingCapture(self.burst_packets(), clock, ticks=0)
|
|
)
|
|
|
|
self.assertIn("Gear captured", printed)
|
|
self.assertEqual(
|
|
reported_pending, [0], "nothing pending: this is the post-loop fallback"
|
|
)
|
|
self.assertIsNotNone(result["gear_path"])
|
|
|
|
def test_an_idle_flow_is_not_decoded_twice_for_the_same_bytes(self):
|
|
# 40 ticks against a flow that never yields gear: without the
|
|
# already-attempted guard this would decode once per tick.
|
|
junk = [
|
|
ParsedIpPacket(
|
|
src_ip=REMOTE_IP,
|
|
dst_ip=LOCAL_IP,
|
|
src_port=30166,
|
|
dst_port=49610,
|
|
payload=b"\x00" * 64,
|
|
protocol="udp",
|
|
)
|
|
]
|
|
_, _, decodes, _ = self.run_capture(
|
|
lambda clock: _TickingCapture(junk, clock, ticks=40)
|
|
)
|
|
|
|
# One live attempt for the single flow size, plus the post-loop scan.
|
|
self.assertEqual(len(decodes), 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|