Files
HyperDRM/tests/flatbuffer-proxy.py

50 lines
1.8 KiB
Python

"""Test-only FlatBuffer TCP relay; accepts only the configured TV address."""
import argparse
import select
import socket
import socketserver
parser = argparse.ArgumentParser()
parser.add_argument('--bind', required=True)
parser.add_argument('--port', type=int, default=19446)
parser.add_argument('--upstream', required=True)
parser.add_argument('--upstream-port', type=int, default=19400)
args = parser.parse_args()
class Relay(socketserver.BaseRequestHandler):
def handle(self):
if self.client_address[0] != args.upstream:
return
transferred = 0
try:
with socket.create_connection((args.upstream, args.upstream_port), timeout=3) as remote:
self.request.settimeout(3)
while True:
readable, _, _ = select.select([self.request, remote], [], [], 3)
if not readable:
break
for source in readable:
data = source.recv(65536)
if not data:
return
destination = remote if source is self.request else self.request
destination.sendall(data)
if source is self.request:
if not transferred:
print('PicCap FlatBuffer data received', flush=True)
transferred += len(data)
except OSError:
pass
finally:
print(f'FlatBuffer connection closed: {transferred} bytes from TV', flush=True)
class Server(socketserver.ThreadingTCPServer):
daemon_threads = True
with Server((args.bind, args.port), Relay) as server:
print('FlatBuffer relay ready', flush=True)
server.serve_forever()