<commit_before>
from nintendo.common import crypto, socket, websocket, scheduler
from nintendo.nex import streams, kerberos

import itertools
import hashlib
import hmac
import struct
import random
import secrets
import zlib
import time

import logging
logger = logging.getLogger(__name__)

TYPE_SYN = 0
TYPE_CONNECT = 1
TYPE_DATA = 2
TYPE_DISCONNECT = 3
TYPE_PING = 4

FLAG_ACK = 1
FLAG_RELIABLE = 2
FLAG_NEED_ACK = 4
FLAG_HAS_SIZE = 8
FLAG_MULTI_ACK = 0x200

OPTION_SUPPORT = 0
OPTION_CONNECTION_SIG = 1
OPTION_FRAGMENT = 2
OPTION_3 = 3
OPTION_4 = 4
OPTION_CONNECTION_SIG_LITE = 0x80


def decode_options(data):
	pos = 0
	options = {}
	while pos < len(data):
		if len(data) - pos < 2:
			logger.error("(Opt) Only one byte left in option data")
			return

		type = data[pos]
		length = data[pos + 1]
		pos += 2

		if len(data) - pos < length:
			logger.error("(Opt) Option length is greater than data size")
			return

		if type == OPTION_SUPPORT:
			if length != 4:
				logger.error("(Opt) Invalid option length in OPTION_SUPPORT")
				return
			options[type] = struct.unpack_from("<I", data, pos)[0]

		elif type in [OPTION_CONNECTION_SIG, OPTION_CONNECTION_SIG_LITE]:
			if length != 16:
				logger.error("(Opt) Invalid option length in OPTION_CONNECTION_SIG")
				return
			options[type] = data[pos : pos + length]

		elif type in [OPTION_FRAGMENT, OPTION_4]:
			if length != 1:
				logger.error("(Opt) Invalid option length in %s",
				             "OPTION_FRAGMENT" if type == OPTION_FRAGMENT else "OPTION_4")
				return
			options[type] = data[pos]

		elif type == OPTION_3:
			if length != 2:
				logger.error("(Opt) Invalid option length in OPTION_3")
				return
			options[type] = struct.unpack_from("<H", data, pos)[0]

		else:
			logger.error("(Opt) Unrecognized option type %i", type)
			return
		
		pos += length
	return options


class PRUDPPacket:
	def __init__(self, type=None, flags=None):
		self.type = type
		self.flags = flags
		
		self.source_port = None
		self.source_type = None
		self.dest_port = None
		self.dest_type = None
		self.session_id = None
		self.packet_id = None
		self.signature = None
		self.fragment_id = 0
		self.multi_ack_version = 0
		self.payload = b""
		
	def __repr__(self):
		return "<PRUDPPacket type=%i flags=%03X seq=%s frag=%s>" %(self.type, self.flags, self.packet_id, self.fragment_id)
	
	
class PRUDPMessageV0:
	def __init__(self, client, settings):
		self.client = client
		self.signature_version = settings.get("prudp_v0.signature_version")
		self.flags_version = settings.get("prudp_v0.flags_version")
		self.checksum_version = settings.get("prudp_v0.checksum_version")
		self.reset()
		
	def reset(self):
		self.buffer = b""
	def signature_size(self): return 4
	
	def checksum_size(self):
		if self.checksum_version == 0:
			return 4
		return 1
		
	def calc_checksum(self, data):
		if self.checksum_version == 0:
			data = data.ljust((len(data) + 3) & ~3, b"\0")
			base = self.client.signature_base & 0xFF
			words = struct.unpack("<%iI" %(len(data) // 4), data)
			return (base + sum(words)) & 0xFFFFFFFF

		else:
			words = struct.unpack_from("<%iI" %(len(data) // 4), data)
			temp = sum(words) & 0xFFFFFFFF
			
			checksum = self.client.signature_base
			checksum += sum(data[len(data) & ~3:])
			checksum += sum(struct.pack("<I", temp))
			return checksum & 0xFF
		
	def calc_data_signature(self, packet):
		data = packet.payload
		if self.signature_version == 0:
			data = self.client.session_key + struct.pack("<HB", packet.packet_id, packet.fragment_id) + data

		if data:
			return hmac.new(self.client.signature_key, data).digest()[:4]
		return struct.pack("<I", 0x12345678)
		
	def calc_packet_signature(self, packet, signature):
		if packet.type == TYPE_DATA: return self.calc_data_signature(packet)
		if packet.type == TYPE_DISCONNECT and self.signature_version == 0:
			return self.calc_data_signature(packet)
		return signature if signature else bytes(4)
		
	def header_format(self):
		if self.flags_version == 0:
			return "<BBBB4sH"
		return "<BBHB4sH"
		
	def encode(self, packet):
		if self.flags_version == 0:
			type_field = packet.type | (packet.flags << 3)
		else:
			type_field = packet.type | (packet.flags << 4)
		
		header = struct.pack(self.header_format(),
			packet.source_port | (packet.source_type << 4),
			packet.dest_port | (packet.dest_type << 4),
			packet.type | (packet.flags << 4),
			packet.session_id,
			self.calc_packet_signature(packet, self.client.target_signature),
			packet.packet_id
		)
		
		options = self.encode_options(packet)
		data = header + options + packet.payload
		if self.checksum_size() == 1:
			data += struct.pack("<B", self.calc_checksum(data))
		elif self.checksum_size() == 4:
			data += struct.pack("<I", self.calc_checksum(data))
		return data
		
	def encode_options(self, packet):
		options = b""
		if packet.type in [TYPE_SYN, TYPE_CONNECT]:
			options += packet.signature
		if packet.type == TYPE_DATA:
			options += struct.pack("<B", packet.fragment_id)
		if packet.flags & FLAG_HAS_SIZE:
			options += struct.pack("<H", len(packet.payload))
		return options
		
	def decode(self, data):
		self.buffer += data
	
		packets = []
		while self.buffer:
			if len(self.buffer) < 12: return packets
			
			#Extract packet header
			source, dest, type_flags, session_id, signature, packet_id = \
				struct.unpack_from(self.header_format(), self.buffer)
				
			packet = PRUDPPacket()
			packet.source_type = source >> 4
			packet.source_port = source & 0xF
			packet.dest_type = dest >> 4
			packet.dest_port = dest & 0xF
			packet.session_id = session_id
			packet.packet_id = packet_id

			if self.flags_version == 0:
				packet.flags = type_flags >> 3
				packet.type = type_flags & 7
			else:
				packet.flags = type_flags >> 4
				packet.type = type_flags & 0xF

			offset = 11
			
			#Extract connection signature
			if packet.type in [TYPE_SYN, TYPE_CONNECT]:
				if len(self.buffer) < offset + 4: return packets
				packet.signature = self.buffer[offset : offset + 4]
				offset += 4

			#Extract fragment id
			if packet.type == TYPE_DATA:
				if len(self.buffer) < offset + 1: return packets
				packet.fragment_id = self.buffer[offset]
				offset += 1

			#Extract payload size
			if packet.flags & FLAG_HAS_SIZE:
				if len(self.buffer) < offset + 2: return packets
				payload_size = struct.unpack_from("<H", self.buffer, offset)[0]
				offset += 2
			else:
				payload_size = len(self.buffer) - offset - self.checksum_size()

			if len(self.buffer) < offset + payload_size + self.checksum_size():
				return packets

			#Extract checksum
			if self.checksum_size() == 1:
				checksum = self.buffer[offset + payload_size]
			elif self.checksum_size() == 4:
				checksum = struct.unpack_from("<I", self.buffer, offset + payload_size)

			#Verify checksum
			expected_checksum = self.calc_checksum(self.buffer[:offset + payload_size])
			if checksum != expected_checksum:
				logger.error("(V0) Invalid checksum (expected %i, got %i)", checksum, expected_checksum)
				self.reset()
				return packets
			
			#Extract payload
			packet.payload = self.buffer[offset : offset + payload_size]
			
			#Verify packet signature
			expected_signature = self.calc_packet_signature(packet, self.client.source_signature)
			if signature != expected_signature:
				logger.error(
					"(V0) Invalid packet signature (expected %s, got %s)",
					signature.hex(), expected_signature.hex()
				)
				self.reset()
				return packets
			
			self.buffer = self.buffer[offset + payload_size + self.checksum_size():]
			packets.append(packet)
		return packets

	
class PRUDPMessageV1:
	def __init__(self, client, settings):
		self.client = client
		self.settings = settings
		self.reset()
		
	def reset(self):
		self.buffer = b""
	def signature_size(self): return 16

	def calc_packet_signature(self, session_key, header, options, signature, payload):
		mac = hmac.new(self.client.signature_key)
		mac.update(header[4:])
		mac.update(session_key)
		mac.update(struct.pack("<I", self.client.signature_base))
		mac.update(signature)
		mac.update(options)
		mac.update(payload)
		return mac.digest()

	def encode(self, packet):
		# On the server side, the session_key has already been set by validate_connection_request
		# the client however still needs an ACK without a session key to validate the packet
		session_key = self.client.session_key
		if packet.type == TYPE_CONNECT and self.client.session_key:
			session_key = b""

		packet.flags |= FLAG_HAS_SIZE
		options = self.encode_options(packet)
		header = self.encode_header(packet, len(options))
		checksum = self.calc_packet_signature(session_key, header, options, self.client.target_signature, packet.payload)
		return b"\xEA\xD0" + header + checksum + options + packet.payload
		
	def encode_header(self, packet, option_size):
		return struct.pack("<BBHBBHBBH",
			1, #PRUDP version
			option_size,
			len(packet.payload),
			packet.source_port | (packet.source_type << 4),
			packet.dest_port | (packet.dest_type << 4),
			packet.type | (packet.flags << 4),
			packet.session_id, packet.multi_ack_version, packet.packet_id
		)
		
	def encode_options(self, packet):
		options = b""
		if packet.type in [TYPE_SYN, TYPE_CONNECT]:
			options += struct.pack("<BBI", OPTION_SUPPORT, 4, self.client.minor_version)
			options += struct.pack("<BB16s", OPTION_CONNECTION_SIG, 16, packet.signature)
			if packet.type == TYPE_CONNECT:
				options += struct.pack("<BBH", OPTION_3, 2, random.randint(0, 0xFFFF))
			options += struct.pack("<BBB", OPTION_4, 1, 0)
		elif packet.type == TYPE_DATA:
			options += struct.pack("<BBB", OPTION_FRAGMENT, 1, packet.fragment_id)
		return options
		
	def decode(self, data):
		self.buffer += data

		packets = []
		while self.buffer:
			if len(self.buffer) < 30: return packets
			if self.buffer[:2] != b"\xEA\xD0":
				logger.error("(V1) Invalid magic number")
				self.reset()
				return packets
		
			packet = PRUDPPacket()

			header = self.buffer[2 : 14]
			checksum = self.buffer[14 : 30]
			
			version, option_size, payload_size, source, dest, type_flags, session_id, \
				packet.multi_ack_version, packet_id = struct.unpack("<BBHBBHBBH", header)

			if version != 1:
				logger.error("(V1) Version check failed")
				self.reset()
				return packets
			if packet.multi_ack_version > 1:
				logger.error("(V1) Unrecognized aggregate ack version: %i" %packet.multi_ack_version)
				self.reset()
				return packets
			if len(self.buffer) < 30 + option_size + payload_size:
				return packets
				
			packet.source_type = source >> 4
			packet.source_port = source & 0xF
			packet.dest_type = dest >> 4
			packet.dest_port = dest & 0xF
			packet.flags = type_flags >> 4
			packet.type = type_flags & 0xF
			packet.session_id = session_id
			packet.packet_id = packet_id
			
			option_data = self.buffer[30 : 30 + option_size:]
			options = decode_options(option_data)
			if options is None:
				self.reset()
				return packets
			
			if packet.type in [TYPE_SYN, TYPE_CONNECT]:
				if OPTION_CONNECTION_SIG not in options:
					logger.error("(V1) Expected connection signature in SYN/CONNECT packet")
					self.reset()
					return packets
				packet.signature = options[OPTION_CONNECTION_SIG]
			elif packet.type == TYPE_DATA:
				if OPTION_FRAGMENT not in options:
					logger.error("(V1) Expected fragment id in DATA packet")
					self.reset()
					return packets
				packet.fragment_id = options[OPTION_FRAGMENT]
			
			packet.payload = self.buffer[30 + option_size : 30 + option_size + payload_size]
			
			if self.calc_packet_signature(self.client.session_key, header, option_data, self.client.source_signature, packet.payload) != checksum:
				logger.error("(V1) Invalid packet signature")
				self.reset()
				return packets
			
			self.buffer = self.buffer[30 + option_size + payload_size:]
			packets.append(packet)
		return packets

		
class PRUDPLiteMessage:
	def __init__(self, client, settings):
		self.client = client
		self.settings = settings
		self.reset()
		
	def reset(self):
		self.buffer = b""
	def signature_size(self): return 16
		
	def encode(self, packet):
		options = self.encode_options(packet)
		header = self.encode_header(packet, len(options))
		return header + options + packet.payload
		
	def encode_header(self, packet, option_size):
		return struct.pack("<BBHBBBBHH",
			0x80, option_size,
			len(packet.payload),
			(packet.source_type << 4) | packet.dest_type,
			packet.source_port,
			packet.dest_port,
			packet.fragment_id,
			packet.type | (packet.flags << 4),
			packet.packet_id
		)
		
	def encode_options(self, packet):
		options = b""
		if packet.type in [TYPE_SYN, TYPE_CONNECT]:
			options += struct.pack("<BBI", OPTION_SUPPORT, 4, self.client.minor_version)
		if packet.type == TYPE_SYN and packet.flags & FLAG_ACK:
			options += struct.pack("<BB16s", OPTION_CONNECTION_SIG, 16, packet.signature)
		if packet.type == TYPE_CONNECT and not packet.flags & FLAG_ACK:
			options += struct.pack("<BB16s", OPTION_CONNECTION_SIG_LITE, 16, packet.signature)
		return options
		
	def decode(self, data):
		self.buffer += data

		packets = []
		while self.buffer:
			if len(self.buffer) < 12: return packets
		
			packet = PRUDPPacket()
			
			magic, option_size, payload_size, stream_types, source_port, dest_port, \
				fragment_id, type_flags, packet_id = struct.unpack_from("<BBHBBBBHH", self.buffer)
			
			if magic != 0x80:
				logger.error("(Lite) Invalid magic number")
				self.reset()
				return packets
			if len(self.buffer) < 12 + option_size + payload_size:
				return packets
				
			packet.source_type = stream_types >> 4
			packet.dest_type = stream_types & 0xF
			packet.source_port = source_port
			packet.dest_port = dest_port
			packet.fragment_id = fragment_id
			packet.flags = type_flags >> 4
			packet.type = type_flags & 0xF
			packet.packet_id = packet_id
			packet.session_id = 0
			
			option_data = self.buffer[12 : 12 + option_size:]
			options = decode_options(option_data)
			if options is None:
				self.reset()
				return packets
			
			if packet.type == TYPE_SYN and packet.flags & FLAG_ACK:
				if OPTION_CONNECTION_SIG not in options:
					logger.error("(Lite) Expected connection signature in SYN/ACK packet")
					self.reset()
					return packets
				packet.signature = options[OPTION_CONNECTION_SIG]
			
			packet.payload = self.buffer[12 + option_size : 12 + option_size + payload_size]
			self.buffer = self.buffer[12 + option_size + payload_size:]
			packets.append(packet)
		return packets

		
class RC4Encryption:
	def __init__(self, key):
		self.rc4enc = crypto.RC4(key)
		self.rc4dec = crypto.RC4(key)
		
	def set_key(self, key):
		self.rc4enc.set_key(key)
		self.rc4dec.set_key(key)
		
	def encrypt(self, data): return self.rc4enc.crypt(data)
	def decrypt(self, data): return self.rc4dec.crypt(data)
	

class DummyEncryption:
	def set_key(self, key): pass
	def encrypt(self, data): return data
	def decrypt(self, data): return data
	
	
class ZlibCompression:
	def compress(self, data):
		compressed = zlib.compress(data)
		ratio = int(len(data) / len(compressed) + 1)
		return bytes([ratio]) + compressed
		
	def decompress(self, data):
		decompressed = zlib.decompress(data[1:])
		ratio = int(len(decompressed) / (len(data) - 1) + 1)
		if ratio != data[0]:
			logger.warning(
				"Unexpected compression ratio (expected %i, got %i)",
				ratio, data[0]
			)
		return decompressed
		
		
class DummyCompression:
	def compress(self, data): return data
	def decompress(self, data): return data
	

class PRUDPClient:

	READY = 0
	ACCEPTING = 1
	CONNECTING = 2
	CONNECTED = 3
	DISCONNECTING = 4
	DISCONNECTED = 5

	DEFAULT_KEY = b"CD&ML"

	def __init__(self, settings, sock=None):
		self.settings = settings
		self.transport_type = settings.get("prudp.transport")
		self.stream_type = settings.get("prudp.stream_type")
		self.fragment_size = settings.get("prudp.fragment_size")
		self.resend_timeout = settings.get("prudp.resend_timeout")
		self.resend_limit = settings.get("prudp.resend_limit")
		self.ping_timeout = settings.get("prudp.ping_timeout")
		self.silence_timeout = settings.get("prudp.silence_timeout")
		self.use_compression = settings.get("prudp.compression")
		self.minor_version = settings.get("prudp.minor_version")
		
		self.sock = sock
		if not self.sock:
			if self.transport_type == settings.TRANSPORT_UDP:
				self.sock = socket.Socket(socket.TYPE_UDP)
			elif self.transport_type == settings.TRANSPORT_TCP:
				self.sock = socket.Socket(socket.TYPE_TCP)
			else:
				self.sock = websocket.WebSocketClient(True)

		access_key = settings.get("server.access_key")
		self.signature_key = hashlib.md5(access_key).digest()
		self.signature_base = sum(access_key)
		
		if self.transport_type == settings.TRANSPORT_UDP:
			if settings.get("prudp.version") == 0:
				self.packet_encoder = PRUDPMessageV0(self, settings)
			else:
				self.packet_encoder = PRUDPMessageV1(self, settings)
			self.encryption = RC4Encryption(self.DEFAULT_KEY)
		else:
			self.packet_encoder = PRUDPLiteMessage(self, settings)
			self.encryption = DummyEncryption()
			
		if settings.get("prudp.compression") == 0:
			self.compression = DummyCompression()
		else:
			self.compression = ZlibCompression()
		
		self.encryption.set_key(self.DEFAULT_KEY)
		self.session_key = b""
			
		self.source_signature = b""
		self.target_signature = b""
		
		self.packets = []
		self.packet_queue = {}
		self.fragment_buffer = b""
		self.packet_id_out = itertools.count()
		self.packet_id_in = 0
		self.local_session_id = 0
		self.remote_session_id = 0
		
		self.ack_events = {}
		self.ping_event = None
		self.timeout_event = None
		self.socket_event = None
		
		self.state = self.READY
		
	def set_session_key(self, key):
		self.encryption.set_key(key)
		self.session_key = key
		
	def is_connected(self): return self.state == self.CONNECTED
	def local_address(self): return self.sock.local_address()
	def remote_address(self): return self.sock.remote_address()
	
	def build_connection_request(self): return b""
	def build_connection_response(self): return b""
	def validate_connection_request(self, payload): return True
	def validate_connection_response(self, payload): return True
		
	def connect(self, host, port, stream_id):
		if self.state != self.READY:
			raise RuntimeError("PRUDP socket may only be used once")
	
		logger.info("Connecting to %s:%i:%i", host, port, stream_id)
		self.state = self.CONNECTING
		
		self.remote_port = stream_id
		
		self.local_port = 0xF<commit_msg>Improve nintendo.nex.prudp class init<commit_after>
from nintendo.common import crypto, socket, websocket, scheduler
from nintendo.nex import streams, kerberos

import itertools
import hashlib
import hmac
import struct
import random
import secrets
import zlib
import time

import logging
logger = logging.getLogger(__name__)

TYPE_SYN = 0
TYPE_CONNECT = 1
TYPE_DATA = 2
TYPE_DISCONNECT = 3
TYPE_PING = 4

FLAG_ACK = 1
FLAG_RELIABLE = 2
FLAG_NEED_ACK = 4
FLAG_HAS_SIZE = 8
FLAG_MULTI_ACK = 0x200

OPTION_SUPPORT = 0
OPTION_CONNECTION_SIG = 1
OPTION_FRAGMENT = 2
OPTION_3 = 3
OPTION_4 = 4
OPTION_CONNECTION_SIG_LITE = 0x80


def decode_options(data):
	pos = 0
	options = {}
	while pos < len(data):
		if len(data) - pos < 2:
			logger.error("(Opt) Only one byte left in option data")
			return

		type = data[pos]
		length = data[pos + 1]
		pos += 2

		if len(data) - pos < length:
			logger.error("(Opt) Option length is greater than data size")
			return

		if type == OPTION_SUPPORT:
			if length != 4:
				logger.error("(Opt) Invalid option length in OPTION_SUPPORT")
				return
			options[type] = struct.unpack_from("<I", data, pos)[0]

		elif type in [OPTION_CONNECTION_SIG, OPTION_CONNECTION_SIG_LITE]:
			if length != 16:
				logger.error("(Opt) Invalid option length in OPTION_CONNECTION_SIG")
				return
			options[type] = data[pos : pos + length]

		elif type in [OPTION_FRAGMENT, OPTION_4]:
			if length != 1:
				logger.error("(Opt) Invalid option length in %s",
				             "OPTION_FRAGMENT" if type == OPTION_FRAGMENT else "OPTION_4")
				return
			options[type] = data[pos]

		elif type == OPTION_3:
			if length != 2:
				logger.error("(Opt) Invalid option length in OPTION_3")
				return
			options[type] = struct.unpack_from("<H", data, pos)[0]

		else:
			logger.error("(Opt) Unrecognized option type %i", type)
			return
		
		pos += length
	return options


class PRUDPPacket:
	def __init__(self, type=None, flags=None):
		self.type = type
		self.flags = flags
		
		self.source_port = None
		self.source_type = None
		self.dest_port = None
		self.dest_type = None
		self.session_id = None
		self.packet_id = None
		self.signature = None
		self.fragment_id = 0
		self.multi_ack_version = 0
		self.payload = b""
		
	def __repr__(self):
		return "<PRUDPPacket type=%i flags=%03X seq=%s frag=%s>" %(self.type, self.flags, self.packet_id, self.fragment_id)
	
	
class PRUDPMessageV0:
	def __init__(self, client, settings):
		self.client = client
		self.signature_version = settings.get("prudp_v0.signature_version")
		self.flags_version = settings.get("prudp_v0.flags_version")
		self.checksum_version = settings.get("prudp_v0.checksum_version")
		self.reset()
		
	def reset(self):
		self.buffer = b""
	def signature_size(self): return 4
	
	def checksum_size(self):
		if self.checksum_version == 0:
			return 4
		return 1
		
	def calc_checksum(self, data):
		if self.checksum_version == 0:
			data = data.ljust((len(data) + 3) & ~3, b"\0")
			base = self.client.signature_base & 0xFF
			words = struct.unpack("<%iI" %(len(data) // 4), data)
			return (base + sum(words)) & 0xFFFFFFFF

		else:
			words = struct.unpack_from("<%iI" %(len(data) // 4), data)
			temp = sum(words) & 0xFFFFFFFF
			
			checksum = self.client.signature_base
			checksum += sum(data[len(data) & ~3:])
			checksum += sum(struct.pack("<I", temp))
			return checksum & 0xFF
		
	def calc_data_signature(self, packet):
		data = packet.payload
		if self.signature_version == 0:
			data = self.client.session_key + struct.pack("<HB", packet.packet_id, packet.fragment_id) + data

		if data:
			return hmac.new(self.client.signature_key, data).digest()[:4]
		return struct.pack("<I", 0x12345678)
		
	def calc_packet_signature(self, packet, signature):
		if packet.type == TYPE_DATA: return self.calc_data_signature(packet)
		if packet.type == TYPE_DISCONNECT and self.signature_version == 0:
			return self.calc_data_signature(packet)
		return signature if signature else bytes(4)
		
	def header_format(self):
		if self.flags_version == 0:
			return "<BBBB4sH"
		return "<BBHB4sH"
		
	def encode(self, packet):
		if self.flags_version == 0:
			type_field = packet.type | (packet.flags << 3)
		else:
			type_field = packet.type | (packet.flags << 4)
		
		header = struct.pack(self.header_format(),
			packet.source_port | (packet.source_type << 4),
			packet.dest_port | (packet.dest_type << 4),
			packet.type | (packet.flags << 4),
			packet.session_id,
			self.calc_packet_signature(packet, self.client.target_signature),
			packet.packet_id
		)
		
		options = self.encode_options(packet)
		data = header + options + packet.payload
		if self.checksum_size() == 1:
			data += struct.pack("<B", self.calc_checksum(data))
		elif self.checksum_size() == 4:
			data += struct.pack("<I", self.calc_checksum(data))
		return data
		
	def encode_options(self, packet):
		options = b""
		if packet.type in [TYPE_SYN, TYPE_CONNECT]:
			options += packet.signature
		if packet.type == TYPE_DATA:
			options += struct.pack("<B", packet.fragment_id)
		if packet.flags & FLAG_HAS_SIZE:
			options += struct.pack("<H", len(packet.payload))
		return options
		
	def decode(self, data):
		self.buffer += data
	
		packets = []
		while self.buffer:
			if len(self.buffer) < 12: return packets
			
			#Extract packet header
			source, dest, type_flags, session_id, signature, packet_id = \
				struct.unpack_from(self.header_format(), self.buffer)
				
			packet = PRUDPPacket()
			packet.source_type = source >> 4
			packet.source_port = source & 0xF
			packet.dest_type = dest >> 4
			packet.dest_port = dest & 0xF
			packet.session_id = session_id
			packet.packet_id = packet_id

			if self.flags_version == 0:
				packet.flags = type_flags >> 3
				packet.type = type_flags & 7
			else:
				packet.flags = type_flags >> 4
				packet.type = type_flags & 0xF

			offset = 11
			
			#Extract connection signature
			if packet.type in [TYPE_SYN, TYPE_CONNECT]:
				if len(self.buffer) < offset + 4: return packets
				packet.signature = self.buffer[offset : offset + 4]
				offset += 4

			#Extract fragment id
			if packet.type == TYPE_DATA:
				if len(self.buffer) < offset + 1: return packets
				packet.fragment_id = self.buffer[offset]
				offset += 1

			#Extract payload size
			if packet.flags & FLAG_HAS_SIZE:
				if len(self.buffer) < offset + 2: return packets
				payload_size = struct.unpack_from("<H", self.buffer, offset)[0]
				offset += 2
			else:
				payload_size = len(self.buffer) - offset - self.checksum_size()

			if len(self.buffer) < offset + payload_size + self.checksum_size():
				return packets

			#Extract checksum
			if self.checksum_size() == 1:
				checksum = self.buffer[offset + payload_size]
			elif self.checksum_size() == 4:
				checksum = struct.unpack_from("<I", self.buffer, offset + payload_size)

			#Verify checksum
			expected_checksum = self.calc_checksum(self.buffer[:offset + payload_size])
			if checksum != expected_checksum:
				logger.error("(V0) Invalid checksum (expected %i, got %i)", checksum, expected_checksum)
				self.reset()
				return packets
			
			#Extract payload
			packet.payload = self.buffer[offset : offset + payload_size]
			
			#Verify packet signature
			expected_signature = self.calc_packet_signature(packet, self.client.source_signature)
			if signature != expected_signature:
				logger.error(
					"(V0) Invalid packet signature (expected %s, got %s)",
					signature.hex(), expected_signature.hex()
				)
				self.reset()
				return packets
			
			self.buffer = self.buffer[offset + payload_size + self.checksum_size():]
			packets.append(packet)
		return packets

	
class PRUDPMessageV1:
	def __init__(self, client, settings):
		self.client = client
		self.settings = settings
		self.reset()
		
	def reset(self):
		self.buffer = b""
	def signature_size(self): return 16

	def calc_packet_signature(self, session_key, header, options, signature, payload):
		mac = hmac.new(self.client.signature_key)
		mac.update(header[4:])
		mac.update(session_key)
		mac.update(struct.pack("<I", self.client.signature_base))
		mac.update(signature)
		mac.update(options)
		mac.update(payload)
		return mac.digest()

	def encode(self, packet):
		# On the server side, the session_key has already been set by validate_connection_request
		# the client however still needs an ACK without a session key to validate the packet
		session_key = self.client.session_key
		if packet.type == TYPE_CONNECT and self.client.session_key:
			session_key = b""

		packet.flags |= FLAG_HAS_SIZE
		options = self.encode_options(packet)
		header = self.encode_header(packet, len(options))
		checksum = self.calc_packet_signature(session_key, header, options, self.client.target_signature, packet.payload)
		return b"\xEA\xD0" + header + checksum + options + packet.payload
		
	def encode_header(self, packet, option_size):
		return struct.pack("<BBHBBHBBH",
			1, #PRUDP version
			option_size,
			len(packet.payload),
			packet.source_port | (packet.source_type << 4),
			packet.dest_port | (packet.dest_type << 4),
			packet.type | (packet.flags << 4),
			packet.session_id, packet.multi_ack_version, packet.packet_id
		)
		
	def encode_options(self, packet):
		options = b""
		if packet.type in [TYPE_SYN, TYPE_CONNECT]:
			options += struct.pack("<BBI", OPTION_SUPPORT, 4, self.client.minor_version)
			options += struct.pack("<BB16s", OPTION_CONNECTION_SIG, 16, packet.signature)
			if packet.type == TYPE_CONNECT:
				options += struct.pack("<BBH", OPTION_3, 2, random.randint(0, 0xFFFF))
			options += struct.pack("<BBB", OPTION_4, 1, 0)
		elif packet.type == TYPE_DATA:
			options += struct.pack("<BBB", OPTION_FRAGMENT, 1, packet.fragment_id)
		return options
		
	def decode(self, data):
		self.buffer += data

		packets = []
		while self.buffer:
			if len(self.buffer) < 30: return packets
			if self.buffer[:2] != b"\xEA\xD0":
				logger.error("(V1) Invalid magic number")
				self.reset()
				return packets
		
			packet = PRUDPPacket()

			header = self.buffer[2 : 14]
			checksum = self.buffer[14 : 30]
			
			version, option_size, payload_size, source, dest, type_flags, session_id, \
				packet.multi_ack_version, packet_id = struct.unpack("<BBHBBHBBH", header)

			if version != 1:
				logger.error("(V1) Version check failed")
				self.reset()
				return packets
			if packet.multi_ack_version > 1:
				logger.error("(V1) Unrecognized aggregate ack version: %i" %packet.multi_ack_version)
				self.reset()
				return packets
			if len(self.buffer) < 30 + option_size + payload_size:
				return packets
				
			packet.source_type = source >> 4
			packet.source_port = source & 0xF
			packet.dest_type = dest >> 4
			packet.dest_port = dest & 0xF
			packet.flags = type_flags >> 4
			packet.type = type_flags & 0xF
			packet.session_id = session_id
			packet.packet_id = packet_id
			
			option_data = self.buffer[30 : 30 + option_size:]
			options = decode_options(option_data)
			if options is None:
				self.reset()
				return packets
			
			if packet.type in [TYPE_SYN, TYPE_CONNECT]:
				if OPTION_CONNECTION_SIG not in options:
					logger.error("(V1) Expected connection signature in SYN/CONNECT packet")
					self.reset()
					return packets
				packet.signature = options[OPTION_CONNECTION_SIG]
			elif packet.type == TYPE_DATA:
				if OPTION_FRAGMENT not in options:
					logger.error("(V1) Expected fragment id in DATA packet")
					self.reset()
					return packets
				packet.fragment_id = options[OPTION_FRAGMENT]
			
			packet.payload = self.buffer[30 + option_size : 30 + option_size + payload_size]
			
			if self.calc_packet_signature(self.client.session_key, header, option_data, self.client.source_signature, packet.payload) != checksum:
				logger.error("(V1) Invalid packet signature")
				self.reset()
				return packets
			
			self.buffer = self.buffer[30 + option_size + payload_size:]
			packets.append(packet)
		return packets

		
class PRUDPLiteMessage:
	def __init__(self, client, settings):
		self.client = client
		self.settings = settings
		self.reset()
		
	def reset(self):
		self.buffer = b""
	def signature_size(self): return 16
		
	def encode(self, packet):
		options = self.encode_options(packet)
		header = self.encode_header(packet, len(options))
		return header + options + packet.payload
		
	def encode_header(self, packet, option_size):
		return struct.pack("<BBHBBBBHH",
			0x80, option_size,
			len(packet.payload),
			(packet.source_type << 4) | packet.dest_type,
			packet.source_port,
			packet.dest_port,
			packet.fragment_id,
			packet.type | (packet.flags << 4),
			packet.packet_id
		)
		
	def encode_options(self, packet):
		options = b""
		if packet.type in [TYPE_SYN, TYPE_CONNECT]:
			options += struct.pack("<BBI", OPTION_SUPPORT, 4, self.client.minor_version)
		if packet.type == TYPE_SYN and packet.flags & FLAG_ACK:
			options += struct.pack("<BB16s", OPTION_CONNECTION_SIG, 16, packet.signature)
		if packet.type == TYPE_CONNECT and not packet.flags & FLAG_ACK:
			options += struct.pack("<BB16s", OPTION_CONNECTION_SIG_LITE, 16, packet.signature)
		return options
		
	def decode(self, data):
		self.buffer += data

		packets = []
		while self.buffer:
			if len(self.buffer) < 12: return packets
		
			packet = PRUDPPacket()
			
			magic, option_size, payload_size, stream_types, source_port, dest_port, \
				fragment_id, type_flags, packet_id = struct.unpack_from("<BBHBBBBHH", self.buffer)
			
			if magic != 0x80:
				logger.error("(Lite) Invalid magic number")
				self.reset()
				return packets
			if len(self.buffer) < 12 + option_size + payload_size:
				return packets
				
			packet.source_type = stream_types >> 4
			packet.dest_type = stream_types & 0xF
			packet.source_port = source_port
			packet.dest_port = dest_port
			packet.fragment_id = fragment_id
			packet.flags = type_flags >> 4
			packet.type = type_flags & 0xF
			packet.packet_id = packet_id
			packet.session_id = 0
			
			option_data = self.buffer[12 : 12 + option_size:]
			options = decode_options(option_data)
			if options is None:
				self.reset()
				return packets
			
			if packet.type == TYPE_SYN and packet.flags & FLAG_ACK:
				if OPTION_CONNECTION_SIG not in options:
					logger.error("(Lite) Expected connection signature in SYN/ACK packet")
					self.reset()
					return packets
				packet.signature = options[OPTION_CONNECTION_SIG]
			
			packet.payload = self.buffer[12 + option_size : 12 + option_size + payload_size]
			self.buffer = self.buffer[12 + option_size + payload_size:]
			packets.append(packet)
		return packets

		
class RC4Encryption:
	def __init__(self, key):
		self.rc4enc = crypto.RC4(key)
		self.rc4dec = crypto.RC4(key)
		
	def set_key(self, key):
		self.rc4enc.set_key(key)
		self.rc4dec.set_key(key)
		
	def encrypt(self, data): return self.rc4enc.crypt(data)
	def decrypt(self, data): return self.rc4dec.crypt(data)
	

class DummyEncryption:
	def set_key(self, key): pass
	def encrypt(self, data): return data
	def decrypt(self, data): return data
	
	
class ZlibCompression:
	def compress(self, data):
		compressed = zlib.compress(data)
		ratio = int(len(data) / len(compressed) + 1)
		return bytes([ratio]) + compressed
		
	def decompress(self, data):
		decompressed = zlib.decompress(data[1:])
		ratio = int(len(decompressed) / (len(data) - 1) + 1)
		if ratio != data[0]:
			logger.warning(
				"Unexpected compression ratio (expected %i, got %i)",
				ratio, data[0]
			)
		return decompressed
		
		
class DummyCompression:
	def compress(self, data): return data
	def decompress(self, data): return data
	

class PRUDPClient:

	READY = 0
	ACCEPTING = 1
	CONNECTING = 2
	CONNECTED = 3
	DISCONNECTING = 4
	DISCONNECTED = 5

	DEFAULT_KEY = b"CD&ML"

	def __init__(self, settings, sock=None):
		self.settings = settings
		self.transport_type = settings.get("prudp.transport")
		self.stream_type = settings.get("prudp.stream_type")
		self.fragment_size = settings.get("prudp.fragment_size")
		self.resend_timeout = settings.get("prudp.resend_timeout")
		self.resend_limit = settings.get("prudp.resend_limit")
		self.ping_timeout = settings.get("prudp.ping_timeout")
		self.silence_timeout = settings.get("prudp.silence_timeout")
		self.use_compression = settings.get("prudp.compression")
		self.minor_version = settings.get("prudp.minor_version")
		
		self.sock = sock
		if not self.sock:
			if self.transport_type == settings.TRANSPORT_UDP:
				self.sock = socket.Socket(socket.TYPE_UDP)
			elif self.transport_type == settings.TRANSPORT_TCP:
				self.sock = socket.Socket(socket.TYPE_TCP)
			else:
				self.sock = websocket.WebSocketClient(True)

		access_key = settings.get("server.access_key")
		self.signature_key = hashlib.md5(access_key).digest()
		self.signature_base = sum(access_key)
		
		if self.transport_type == settings.TRANSPORT_UDP:
			if settings.get("prudp.version") == 0:
				self.packet_encoder = PRUDPMessageV0(self, settings)
			else:
				self.packet_encoder = PRUDPMessageV1(self, settings)
			self.encryption = RC4Encryption(self.DEFAULT_KEY)
		else:
			self.packet_encoder = PRUDPLiteMessage(self, settings)
			self.encryption = DummyEncryption()
			
		if settings.get("prudp.compression") == 0:
			self.compression = DummyCompression()
		else:
			self.compression = ZlibCompression()
		
		self.encryption.set_key(self.DEFAULT_KEY)
		self.session_key = b""
			
		self.source_signature = b""
		self.target_signature = b""
		
		self.packets = []
		self.packet_queue = {}
		self.fragment_buffer = b""
		self.packet_id_out = itertools.count()
		self.packet_id_in = 0
		self.local_session_id = 0
		self.remote_session_id = 0
		
		self.ack_events = {}
		self.ping_event = None
		self.timeout_event = None
		self.socket_event = None
		self.remote_port = None
		self.local_port = None
		
		self.state = self.READY
		
	def set_session_key(self, key):
		self.encryption.set_key(key)
		self.session_key = key
		
	def is_connected(self): return self.state == self.CONNECTED
	def local_address(self): return self.sock.local_address()
	def remote_address(self): return self.sock.remote_address()
	
	def build_connection_request(self): return b""
	def build_connection_response(self): return b""
	def validate_connection_request(self, payload): return True
	def validate_connection_response(self, payload): return True
		
	def connect(self, host, port, stream_id):
		if self.state != self.READY:
			raise RuntimeError("PRUDP socket may only be used once")
	
		logger.info("Connecting to %s:%i:%i", host, port, stream_id)
		self.state = self.CONNECTING
		
		self.remote_port = stream_id
		
		self.local_port = 0xF
{'fraction_non_alphanumeric': 0.11550628633229687, 'fraction_numerical': 0.01287008246586454, 'mean_word_length': 3.8512591815320043, 'pattern_counts': {'":': 2, '<': 125, '<?xml version=': 0, '>': 43, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 4, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6848424', 'n_tokens_mistral': 14533, 'n_tokens_neox': 12286, 'n_words': 2922}