diff --git a/decode.py b/decode.py index 28e8b97..6fa28d3 100644 --- a/decode.py +++ b/decode.py @@ -25,38 +25,13 @@ args = parser.parse_args() # TODO - probably a better way to do this -def live_capture(): - while True: - # TODO: This is wrong: - # during scapy.all.sniff construction we lose about 27 packets. - # we work around it by capturing 2* packets(frame), so we just lose - # about half the frames. - cap = scapy.all.sniff(iface="enp1s0f0", count=32 * 2) - for pkt in cap: - l = len(bytes(pkt)) - if l != 6972: - continue - data = bytes(pkt) - yield data - - -if args.live: - print('live stream, import scapy') - import scapy.all - print('open stream') - stream = live_capture() - start = datetime.now() - timestamp = start.strftime('%Y%m%d_%H%M%S') - basename = f'live_{timestamp}' - blocks = tqdm(stream) -else: - # Now use args.input_file as the file to process - input_file = args.input_file - basename = os.path.splitext(os.path.basename(input_file))[0] - stream = open(input_file, 'rb') - # Read packets from a pcap file - scanner = pcapng.scanner.FileScanner(stream) - blocks = tqdm(scanner) +def live_capture_cb(cb): + def outer(pkt): + data = bytes(pkt) + l = len(data) + if l == 6972: + cb(data) + scapy.all.sniff(iface="enp1s0f0", filter='udp', prn=outer) def rightsize(it): @@ -77,7 +52,11 @@ def rightsize(it): def removestart(it): "Remove the UDP header from the packets" for x in it: - yield x[0x2A:] + yield removestart_inner(x) + + +def removestart_inner(x): + return x[0x2A:] # Function to parse packet data @@ -97,28 +76,51 @@ def parsed(it): yield parse(x) +class FrameCollector: + def __init__(self): + self.current = [] + + def handle(self, obj): + ret = None + if obj['part'] == 0: + if len(self.current) > 0: + ret = b"".join(self.current) + self.current = [] + #otherdata = [] + self.current.append(obj["data"]) + return ret + #otherdata.append(obj) + + def last(self): + if len(self.current) > 0: + return b"".join(current) + return None + + # Function to group data into frames def frames(it): - current = [] + handler = FrameCollector() #otherdata = [] for obj in it: - if obj['part'] == 0: - if len(current) > 0: - yield b"".join(current) - current = [] - #otherdata = [] - current.append(obj["data"]) - #otherdata.append(obj) - if len(current) > 0: - yield b"".join(current) + ret = handler.handle(obj) + if ret: + yield ret + last = handler.last() + if last: + yield last + WIDTH = 384 HEIGHT = 288 +def bad_frame(frame, width=WIDTH, height=HEIGHT): + return len(frame) != width * height * 2 # 16 bpp + + def skip_bad_frames(it, width=WIDTH, height=HEIGHT): for frame in it: - if len(frame) != width * height * 2: # 16 bpp + if bad_frame(frame): # 16 bpp # Will be fixed when we stopped doing restarts #print(f'{len(frame)} != {width} * {height} * 2') continue @@ -130,16 +132,24 @@ def iterimages(it, width=WIDTH, height=HEIGHT, pixelformat=">H"): yield Image.fromarray(np.frombuffer(frame, dtype=pixelformat).reshape(width, height)) -# Get frames and convert them to images -frames = skip_bad_frames(frames(parsed(removestart(rightsize(blocks))))) +def process_video(): + # Now use args.input_file as the file to process + input_file = args.input_file + basename = os.path.splitext(os.path.basename(input_file))[0] + stream = open(input_file, 'rb') + # Read packets from a pcap file + scanner = pcapng.scanner.FileScanner(stream) + blocks = tqdm(scanner) -# Create the directory for frames if not exists -frame_dir = f"frames/{basename}" -if not os.path.exists(frame_dir): - os.makedirs(frame_dir) + # Get frames and convert them to images + frames = skip_bad_frames(frames(parsed(removestart(rightsize(blocks))))) -# Save each image as a PNG file -if not args.live: + # Create the directory for frames if not exists + frame_dir = f"frames/{basename}" + if not os.path.exists(frame_dir): + os.makedirs(frame_dir) + + # Save each image as a PNG file images = iterimages(it=frames) for i, img in enumerate(images): img.save(f'frames/{basename}/{basename}_{i:04}.png') @@ -162,7 +172,8 @@ if not args.live: subprocess.run(command) print("to play: ffplay thermal.mp4") -else: + +if args.live: # TODO: to video via ffmpeg; right now just a single png # of the last frame def todo_live_ffmpeg(): @@ -175,7 +186,33 @@ else: for frame in frames: fd.write(frame) - for frame in frames: - print('.') + print('live stream, import scapy') + import scapy.all + print('open stream') + + class PacketHandler: + def __init__(self, cb): + self.frame_collector = FrameCollector() + self.cb = cb + + def handle(self, pkt): + pkt = removestart_inner(pkt) + parsed = parse(pkt) + frame_maybe = self.frame_collector.handle(parsed) + if not frame_maybe or bad_frame(frame_maybe): + return + self.cb(frame_maybe) + + progress = tqdm() + + def on_frame(frame): + progress.update(1) Image.fromarray(np.frombuffer(frame, dtype='>H').reshape(WIDTH, HEIGHT)).save(f'live.new.png') os.rename('live.new.png', 'live.png') + + handler = PacketHandler(on_frame) + live_capture_cb(handler.handle) + +else: + process_video() +