From 1a0c63872c63d43de2075f8277eb109cb973e3f4 Mon Sep 17 00:00:00 2001
From: Sunpy <sunpy@osufx.com>
Date: Tue, 19 Mar 2019 17:25:05 +0100
Subject: [PATCH] Deserialize (load)

---
 osuRepy/helpers/typeSerializer.py | 68 +++++++++++++++++++------------
 osuRepy/replay.py                 | 24 ++++++++---
 test.py                           |  4 +-
 3 files changed, 65 insertions(+), 31 deletions(-)

diff --git a/osuRepy/helpers/typeSerializer.py b/osuRepy/helpers/typeSerializer.py
index 81b1aed..1b05fdf 100644
--- a/osuRepy/helpers/typeSerializer.py
+++ b/osuRepy/helpers/typeSerializer.py
@@ -96,14 +96,13 @@ class Serializable:
 		self.prefix = prefix
 		self.filter = filter
 
-	def get_struct_fmt(self, length = 1):
+	def get_struct_fmt(self, length = 1, uleb_length_override = -1):
 		struct_fmt = self.type[0]
 		if length > 1:
 			struct_fmt = b"%d%b" % (length, self.type[0])
 		if self.length is not None:
 			if self.length == TYPE_ULEB:
-				val = self.get_value()
-				struct_fmt = b"%d%b%b" % (len(val[0]), self.length[0], struct_fmt)
+				struct_fmt = b"%d%b%b" % (len(self.get_value()[0]) if uleb_length_override == -1 else uleb_length_override, self.length[0], struct_fmt)
 			else:
 				struct_fmt = b"%b%b" % (self.length[0], struct_fmt)
 		return struct_fmt
@@ -111,7 +110,7 @@ class Serializable:
 	def get_value(self):
 		val = [self.value]
 		if self.filter is not None:
-			val[-1] = self.filter(val[-1])
+			val[-1] = self.filter(val[-1], 0)
 		if self.length is not None:
 			if self.length == TYPE_ULEB:
 				val.insert(0, encode_uleb(len(val[-1])))
@@ -129,10 +128,40 @@ class Serializable:
 			packed = self.prefix + packed
 		return packed
 
-	def unpack(self, length = -1, byte_order = BYTEORDER_LITTLE):
+	def unpack(self, stream, byte_order = BYTEORDER_LITTLE):
+		length = 1
+		read_pos = 0
+		extra_read = 0
+		remove_struct_fmt = []
 		if self.type == TYPE_STRING:
-			length = len(self.value)
-		struct_fmt = byte_order + self.get_struct_fmt(length)
+			if self.length == TYPE_ULEB:
+				extra_read = 1
+				read_pos += 1 # Skip 0x0b
+				length, _read_pos = decode_uleb(stream[read_pos:])
+				read_pos += _read_pos - 1
+			else:
+				l = Serializable(0, self.length)
+				read_pos = l.unpack(stream)
+				length = l.get_value()[-1]
+				remove_struct_fmt.append(0)
+
+
+		struct_fmt = byte_order + self.get_struct_fmt(length, read_pos)
+		length += extra_read
+
+		if len(remove_struct_fmt):
+			for rem in remove_struct_fmt:
+				struct_fmt = struct_fmt[:rem + 1] + struct_fmt[rem + 2:]
+
+		data = struct.unpack(struct_fmt, stream[read_pos:read_pos + self.type[1] * length])
+
+		if self.filter is not None:
+			data = list(data)
+			data[-1] = self.filter(data[-1], 1)
+
+		self.value = data[-1]
+
+		return read_pos + self.type[1] * length
 
 class Serializer:
 	def __init__(self, byte_order = BYTEORDER_LITTLE):
@@ -145,26 +174,15 @@ class Serializer:
 		else:
 			self.stack.append(serializable)
 
+	def serialize(self):
+		return b"".join(s.pack() for s in self.stack)
+
+	def deserialize(self, stream):
+		for s in self.stack:
+			stream = stream[s.unpack(stream):]
+
 	def flush(self):
 		return b"".join(s.pack() for s in self.stack)
 
 	def __bytes__(self):
 		return self.flush()
-
-class Deserializer:
-	def __init__(self, byte_order = BYTEORDER_LITTLE):
-		self.stack = []
-		self.byte_order = byte_order
-
-	def add(self, serializable):
-		if type(serializable) is list:
-			self.stack += serializable
-		else:
-			self.stack.append(serializable)
-	"""
-	def deserialize_data(self, stream):
-		pos = 0
-		for s in self.stack:
-			struct.
-			s.value 
-	"""
\ No newline at end of file
diff --git a/osuRepy/replay.py b/osuRepy/replay.py
index 739cf3c..28f7d28 100644
--- a/osuRepy/replay.py
+++ b/osuRepy/replay.py
@@ -1,6 +1,6 @@
 import struct
 import string
-#import lzma
+import lzma
 import pylzma # Required cause pythons standard library lzma has missing support
 
 from .frame import ReplayFrame
@@ -34,14 +34,21 @@ def md5_file(file):
 		hash.update(chunk)
 	return hash.hexdigest()
 
-def append_and_compress(data):
-	return lzma_compress(data + b"-12345|0|0|1337,")
+def append_and_compress(data, mode):
+	if mode == 0:
+		return lzma_compress(data + b"-12345|0|0|1337,")
+	else:
+		return lzma_decompress(data)[:-len(b"-12345|0|0|1337,")]
 
 def lzma_compress(data):
 	comp = pylzma.compress(data, dictionary = 21, fastBytes = 255, eos = False)
 	comp = comp[:5] + struct.pack(b"<Q", len(data)) + comp[5:] # Append uncompressed length
 	return comp
 
+def lzma_decompress(data):
+	comp = lzma.decompress(data)
+	return comp
+
 class Replay:
 	_mode					= Serializable(osuModes.STANDARD,					TYPE_BYTE)
 	_osu_version			= Serializable(20181216,							TYPE_INT)
@@ -248,10 +255,17 @@ class Replay:
 			serializer.add( self.__getattribute__(attrib) ) # Add value to serializer
 
 		with open(filename, "wb") as f:
-			f.write( serializer.flush() )
+			f.write( serializer.serialize() )
 
 	def load(self, filename):
+		serializer = Serializer()
+
 		attribs = [ a for a in self.__dir__() if not a.startswith("__") ]
 		for attrib in attribs:
 			if attrib == "_end_of_attributes":
-				break
\ No newline at end of file
+				break
+
+			serializer.add( self.__getattribute__(attrib) ) # Add value to serializer
+
+		with open(filename, "rb") as f:
+			serializer.deserialize( f.read() )
diff --git a/test.py b/test.py
index d19cc90..d5892b1 100644
--- a/test.py
+++ b/test.py
@@ -5,4 +5,6 @@ from osuRepy.helpers import osuButtons, osuMods # Enum helpers
 
 replay = Replay() # Create replay instance
 
-replay.load("replay.osr")
\ No newline at end of file
+replay.load("replay.osr")
+
+replay.save("test.osr")
\ No newline at end of file