"""Local binary communication lab (Python 3 standard library).
Run: python binary_transport.py
HTTP: POST http://COMPUTER_IP:8080/binary with a Base64 text body.
WebSocket: ws://COMPUTER_IP:8080/ws, binary messages only.
Replies echo the exact payload. This is a local test service, not a production server.
"""
import base64
import hashlib
import struct
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

LIMIT = 4096


def exact(stream, count):
    data = stream.read(count)
    if len(data) != count:
        raise EOFError
    return data


def packet(opcode, data):
    length = len(data)
    header = bytes([0x80 | opcode, length]) if length < 126 else bytes([0x80 | opcode, 126]) + struct.pack('!H', length)
    return header + data


class Lab(BaseHTTPRequestHandler):
    protocol_version = 'HTTP/1.1'

    def response(self, status, data):
        self.send_response(status)
        self.send_header('Content-Type', 'text/plain; charset=utf-8')
        self.send_header('Content-Length', str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def do_POST(self):
        if self.path != '/binary':
            self.response(404, b'Use /binary')
            return
        try:
            length = int(self.headers.get('Content-Length', '0'))
            if not 0 < length <= LIMIT:
                self.close_connection = True
                self.response(413, b'Invalid body length')
                return
            raw = base64.b64decode(exact(self.rfile, length), validate=True)
        except (ValueError, EOFError):
            self.response(400, b'Invalid Base64')
            return
        print('HTTP received:', raw.hex(' ').upper(), flush=True)
        self.response(200, base64.b64encode(raw))

    def do_GET(self):
        if self.path != '/ws' or self.headers.get('Upgrade', '').lower() != 'websocket':
            self.response(404, b'Use WebSocket /ws or POST /binary')
            return
        key = self.headers.get('Sec-WebSocket-Key', '')
        try:
            if len(base64.b64decode(key, validate=True)) != 16:
                raise ValueError
        except ValueError:
            self.response(400, b'Invalid WebSocket key')
            return
        accept = base64.b64encode(hashlib.sha1((key + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11').encode()).digest()).decode()
        self.send_response(101)
        self.send_header('Upgrade', 'websocket')
        self.send_header('Connection', 'Upgrade')
        self.send_header('Sec-WebSocket-Accept', accept)
        self.end_headers()
        self.close_connection = True
        pending = None
        try:
            while True:
                a, b = exact(self.rfile, 2)
                opcode, final, length = a & 15, bool(a & 128), b & 127
                if length == 126:
                    length = struct.unpack('!H', exact(self.rfile, 2))[0]
                elif length == 127:
                    length = struct.unpack('!Q', exact(self.rfile, 8))[0]
                if a & 112 or not b & 128 or length > LIMIT:
                    self.wfile.write(packet(8, struct.pack('!H', 1002)))
                    return
                if opcode >= 8 and (not final or length > 125):
                    self.wfile.write(packet(8, struct.pack('!H', 1002)))
                    return
                mask = exact(self.rfile, 4)
                raw = exact(self.rfile, length)
                data = bytes(value ^ mask[i % 4] for i, value in enumerate(raw))
                if opcode == 8:
                    self.wfile.write(packet(8, data))
                    return
                if opcode == 9:
                    self.wfile.write(packet(10, data))
                    continue
                if opcode == 10:
                    continue
                if opcode == 2 and pending is None:
                    pending = bytearray(data)
                elif opcode == 0 and pending is not None:
                    pending.extend(data)
                else:
                    self.wfile.write(packet(8, struct.pack('!H', 1003)))
                    return
                if len(pending) > LIMIT:
                    self.wfile.write(packet(8, struct.pack('!H', 1009)))
                    return
                if final:
                    print('WebSocket received:', pending.hex(' ').upper(), flush=True)
                    self.wfile.write(packet(2, bytes(pending)))
                    pending = None
        except (EOFError, OSError):
            return


if __name__ == '__main__':
    print('Binary lab: HTTP /binary and WebSocket /ws on port 8080', flush=True)
    ThreadingHTTPServer(('0.0.0.0', 8080), Lab).serve_forever()