def load(self, data):
sha = data[:32]
if sha != hashlib.sha256(data[32:]).digest():
raise ValueError("Incorrect SHA256 hash")
wup = self.is_wiiu(data)
endian = ">" if wup else "<"
stream = streams.StreamIn(data[32:], endian)
self.title_id = stream.u64()
self.title_version = stream.u32()
self.unk1 = stream.u32()
self.unk2 = stream.u32()
self.unk3 = stream.read(16)
self.unk4 = stream.u32()
self.unk5 = stream.u64()
self.strings = []
for i in range(16):
self.strings.append(IDBEStrings.load(stream))
if wup:
self.tga = stream.read(0x1002C)
self.unk6 = stream.u32()
else:
self.unk6 = stream.read(0x1680)
def test_pushpop(self): stream = streams.StreamIn(b"HelloWorld") stream.seek(5) stream.push() assert stream.read(5) == b"World" stream.pop() assert stream.read(5) == b"World"
def test_available(self): stream = streams.StreamIn(b"HelloWorld") assert stream.available() == 10 stream.skip(5) assert stream.available() == 5 stream.seek(2) assert stream.available() == 8
def test_skip(self): stream = streams.StreamIn(b"HelloWorld") stream.skip(2) assert stream.read(5) == b"lloWo" stream.skip(3) with pytest.raises(OverflowError): stream.skip(1)
def test_align(self): stream = streams.StreamIn(b"HelloWorld") stream.align(100) assert stream.tell() == 0 stream.skip(1) stream.align(4) assert stream.tell() == 4 assert stream.read(5) == b"oWorl"
def test_seek(self): stream = streams.StreamIn(b"HelloWorld") assert stream.read(10) == b"HelloWorld" stream.seek(5) assert stream.read(5) == b"World" stream.seek(10) with pytest.raises(OverflowError): stream.seek(11)
def test_uint(self): stream = streams.StreamIn(b"\x7F\x00\x80\xFF\xFF\xFF\0\0\0\0\xEF\xCD\xAB\x89\x67\x45\x23\x01") assert stream.u8() == 127 assert stream.u16() == 0x8000 assert stream.u24() == 0xFFFFFF assert stream.u32() == 0 assert stream.u64() == 0x123456789ABCDEF stream.skip(-1) with pytest.raises(OverflowError): stream.u16()
def decode(self, data):
self.buffer += data
packets = []
while self.buffer:
if len(self.buffer) < 12: return packets
stream = streams.StreamIn(self.buffer)
if stream.u8() != 0x80:
raise ValueError("(Lite) Invalid magic number")
option_size = stream.u8()
payload_size = stream.u16()
if len(self.buffer) < 12 + option_size + payload_size:
return packets
self.buffer = self.buffer[12 + option_size + payload_size:]
packet = PRUDPPacket()
stream_types = stream.u8()
if stream_types >> 4 != self.remote_stream_type:
raise ValueError("(Lite) Received packet with invalid source stream type")
if stream_types & 0xF != self.local_stream_type:
raise ValueError("(Lite) Received packet with invalid destination stream type")
packet.source_port = stream.u8()
packet.dest_port = stream.u8()
packet.fragment_id = stream.u8()
type_flags = stream.u16()
packet.flags = type_flags >> 4
packet.type = type_flags & 0xF
packet.packet_id = stream.u16()
packet.session_id = 0
option_data = stream.read(option_size)
options = decode_options(option_data)
if not self.verify_options(packet, options):
raise ValueError("(Lite) Received unexpected set of options")
packet.connection_signature = b""
if packet.type in [TYPE_SYN, TYPE_CONNECT]:
packet.minor_version = options[OPTION_SUPPORT] & 0xFF
packet.supported_functions = options[OPTION_SUPPORT] >> 8
if packet.type == TYPE_SYN and packet.flags & FLAG_ACK:
packet.connection_signature = options[OPTION_CONNECTION_SIG]
if packet.type == TYPE_CONNECT and not packet.flags & FLAG_ACK:
packet.signature = options[OPTION_CONNECTION_SIG_LITE]
packet.payload = stream.read(payload_size)
packets.append(packet)
return packets
def decode(self, data):
packets = []
stream = streams.StreamIn(data)
while not stream.eof():
if stream.read(2) != b"\xEA\xD0":
raise ValueError("(V1) Invalid magic number")
header = stream.peek(12)
if stream.u8() != 1:
raise ValueError("(V1) Version check failed")
option_size = stream.u8()
payload_size = stream.u16()
source = stream.u8()
dest = stream.u8()
type_flags = stream.u16()
if source >> 4 != self.remote_stream_type:
raise ValueError("(V1) Received packet with invalid source stream type")
if dest >> 4 != self.local_stream_type:
raise ValueError("(V1) Received packet with invalid destination stream type")
packet = PRUDPPacket()
packet.source_port = source & 0xF
packet.dest_port = dest & 0xF
packet.flags = type_flags >> 4
packet.type = type_flags & 0xF
packet.session_id = stream.u8()
packet.substream_id = stream.u8()
packet.packet_id = stream.u16()
packet.signature = stream.read(16)
option_data = stream.read(option_size)
options = decode_options(option_data)
if not self.verify_options(packet, options):
raise ValueError("(V1) Received unexpected set of options")
if packet.type in [TYPE_SYN, TYPE_CONNECT]:
packet.minor_version = options[OPTION_SUPPORT] & 0xFF
packet.supported_functions = options[OPTION_SUPPORT] >> 8
packet.connection_signature = options[OPTION_CONNECTION_SIG]
packet.max_substream_id = options[OPTION_MAX_SUBSTREAM_ID]
if packet.type == TYPE_CONNECT:
packet.initial_unreliable_id = options[OPTION_UNRELIABLE_SEQ_ID]
if packet.type == TYPE_DATA:
packet.fragment_id = options[OPTION_FRAGMENT_ID]
packet.payload = stream.read(payload_size)
packets.append(packet)
return packets
def decode_options(data):
options = {}
stream = streams.StreamIn(data)
while not stream.eof():
type = stream.u8()
length = stream.u8()
if type not in OPTIONS:
raise ValueError("(Opt) Unrecognized option type: %i" %type)
expected_length, name, format = OPTIONS[type]
if length != expected_length:
raise ValueError("(Opt) Invalid option length in %s" %name)
if type in options:
raise ValueError("(Opt) %s is present more than once" %name)
value = struct.unpack("<" + format, stream.read(length))[0]
options[type] = value
return options
def test_char(self): stream = streams.StreamIn(b"ABC\0D\0") assert stream.char() == "A" assert stream.char() == "B" assert stream.wchar() == "C" assert stream.wchar() == "D"
def test_bool(self): stream = streams.StreamIn(b"\x00\x01\x80") assert stream.bool() is False assert stream.bool() is True assert stream.bool() is True
def test_float(self): stream = streams.StreamIn(b"\0\0\0\x40\0\0\0\0\0\0\0\xC0") assert stream.float() == 2 assert stream.double() == -2
def test_sint(self): stream = streams.StreamIn(b"\x7F\x00\x80\0\0\0\0\xEF\xCD\xAB\x89\x67\x45\x23\x01") assert stream.s8() == 127 assert stream.s16() == -0x8000 assert stream.s32() == 0 assert stream.s64() == 0x123456789ABCDEF
def test_peek(self): stream = streams.StreamIn(b"HelloWorld") assert stream.peek(5) == b"Hello" assert stream.peek(10) == b"HelloWorld" with pytest.raises(OverflowError): stream.peek(11)
def test_size(self): stream = streams.StreamIn(b"HelloWorld") assert stream.size() == 10
def test_read(self): stream = streams.StreamIn(b"HelloWorld") assert stream.read(5) == b"Hello" assert stream.read(5) == b"World" with pytest.raises(OverflowError): stream.read(1)
def test_chars(self): stream = streams.StreamIn(b"ABC\0D\0") assert stream.chars(2) == "AB" assert stream.wchars(2) == "CD"
def test_tell(self): stream = streams.StreamIn(b"HelloWorld") assert stream.tell() == 0 stream.skip(5) assert stream.tell() == 5
def test_repeat(self): stream = streams.StreamIn(b"HelloWorld") assert stream.repeat(stream.char, 5) == ["H", "e", "l", "l", "o"]
def test_eof(self): stream = streams.StreamIn(b"HelloWorld") assert not stream.eof() stream.seek(10) assert stream.eof()
def decode(self, data):
packets = []
stream = streams.StreamIn(data)
while not stream.eof():
start = stream.tell()
source = stream.u8()
dest = stream.u8()
if source >> 4 != self.remote_stream_type:
raise ValueError("(V0) Received packet with invalid source stream type")
if dest >> 4 != self.local_stream_type:
raise ValueError("(V0) Received packet with invalid destination stream type")
packet = PRUDPPacket()
packet.source_port = source & 0xF
packet.dest_port = dest & 0xF
if self.flags_version == 0:
type_flags = stream.u8()
packet.flags = type_flags >> 3
packet.type = type_flags & 7
else:
type_flags = stream.u16()
packet.flags = type_flags >> 4
packet.type = type_flags & 0xF
packet.session_id = stream.u8()
packet.signature = stream.read(4)
packet.packet_id = stream.u16()
if packet.type in [TYPE_SYN, TYPE_CONNECT]:
packet.connection_signature = stream.read(4)
if packet.type == TYPE_DATA:
packet.fragment_id = stream.u8()
if packet.flags & FLAG_HAS_SIZE:
payload_size = stream.u16()
else:
if self.checksum_version == 0:
payload_size = stream.available() - 4
else:
payload_size = stream.available() - 1
packet.payload = stream.read(payload_size)
# Check packet checkusm
end = stream.tell()
checksum_data = stream.get()[start : end]
expected_checksum = self.calc_checksum(checksum_data)
if self.checksum_version == 0:
checksum = stream.u32()
else:
checksum = stream.u8()
if checksum != expected_checksum:
raise ValueError("(V0) Invalid checksum (expected %i, got %i)" %(expected_checksum, checksum))
# Checksum is good!
packets.append(packet)
return packets
def test_readall(self): stream = streams.StreamIn(b"HelloWorld") stream.skip(2) assert stream.readall() == b"lloWorld" assert stream.readall() == b""
def test_pad(self): stream = streams.StreamIn(b"\0\0\0AAABBB") stream.pad(3) stream.pad(3, b"A") with pytest.raises(ValueError): stream.pad(3, b"A")
def test_ascii(self): stream = streams.StreamIn(b"HelloWorld") assert stream.ascii(10) == "HelloWorld"
def parse(cls, data):
instance = cls()
instance.decode(streams.StreamIn(data, ">"))
return instance
def test_get(self): stream = streams.StreamIn(b"HelloWorld") assert stream.get() == b"HelloWorld"
