TypeTreeNode fix - m_Typeflags optional and TPK flat to hierachy · SNWCreations/UnityPy@c694e60 · GitHub
Skip to content

Commit c694e60

Browse files
K0lb3K0lb3
committed
TypeTreeNode fix - m_Typeflags optional and TPK flat to hierachy
1 parent 5e455b0 commit c694e60

5 files changed

Lines changed: 40 additions & 51 deletions

File tree

‎UnityPy/files/ObjectReader.py‎

Lines changed: 2 additions & 2 deletions

‎UnityPy/helpers/Tpk.py‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,16 @@ def init():
1313
import os
1414

1515
with open(
16-
os.path.join(os.path.dirname(os.path.dirname(__file__)), "resources", "uncompressed.tpk"), "rb"
16+
os.path.join(
17+
os.path.dirname(os.path.dirname(__file__)), "resources", "uncompressed.tpk"
18+
),
19+
"rb",
1720
) as f:
1821
global TPKTYPETREE
1922
TPKTYPETREE = TpkFile(f).GetDataBlob()
2023

2124

22-
def get_typetree_nodes(class_id: int, version: tuple):
25+
def get_typetree_node(class_id: int, version: tuple):
2326
global NODES_CACHE
2427
key = (class_id, version)
2528
if key in NODES_CACHE:
@@ -31,12 +34,12 @@ def get_typetree_nodes(class_id: int, version: tuple):
3134
if class_info is None:
3235
raise ValueError("Could not find class info for class id {}".format(class_id))
3336

34-
nodes = generate_flat_nodes(class_info)
35-
NODES_CACHE[key] = nodes
36-
return nodes
37+
node = generate_node(class_info)
38+
NODES_CACHE[key] = node
39+
return node
3740

3841

39-
def generate_flat_nodes(class_info: TpkUnityClass) -> List[TypeTreeNode]:
42+
def generate_node(class_info: TpkUnityClass) -> TypeTreeNode:
4043
nodes = []
4144
NODES = TPKTYPETREE.NodeBuffer.Nodes
4245
stack = [(class_info.ReleaseRootNode, 0)]
@@ -57,7 +60,7 @@ def generate_flat_nodes(class_info: TpkUnityClass) -> List[TypeTreeNode]:
5760
)
5861
stack = [(node_id, level + 1) for node_id in node.SubNodes] + stack
5962
index += 1
60-
return nodes
63+
return TypeTreeNode.from_list(nodes)
6164

6265

6366
######################################################################################
@@ -292,7 +295,9 @@ def fromString(version: str) -> UnityVersion:
292295
return UnityVersion(version.split("."))
293296

294297
@staticmethod
295-
def fromList(major: int = 0, minor: int = 0, patch: int = 0, build: int = 0) -> UnityVersion:
298+
def fromList(
299+
major: int = 0, minor: int = 0, patch: int = 0, build: int = 0
300+
) -> UnityVersion:
296301
return UnityVersion(major << 48 | minor << 32 | patch << 16 | build)
297302

298303
@property

‎UnityPy/helpers/TypeTreeNode.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,9 @@ class TypeTreeNodeC:
2020
m_Type: str
2121
m_Name: str
2222
m_ByteSize: int
23-
m_TypeFlags: int
2423
m_Version: int
2524
m_Children: List[TypeTreeNode] = field(factory=list)
25+
m_TypeFlags: Optional[int] = None
2626
m_VariableCount: Optional[int] = None
2727
m_Index: Optional[int] = None
2828
m_MetaFlag: Optional[int] = None
@@ -92,7 +92,7 @@ def read_string(reader: EndianBinaryReader, value: int) -> str:
9292
offset = value & 0x7FFFFFFF
9393
return CommonString.get(offset, str(offset))
9494

95-
fake_root: TypeTreeNode = cls(-1, "", "", 0, 0, 0, [])
95+
fake_root: TypeTreeNode = cls(-1, "", "", 0, 0, [])
9696
stack: List[TypeTreeNode] = [fake_root]
9797
parent = fake_root
9898
prev = fake_root
@@ -119,13 +119,14 @@ def read_string(reader: EndianBinaryReader, value: int) -> str:
119119

120120
@classmethod
121121
def from_list(cls, nodes: List[dict]) -> TypeTreeNode:
122-
fake_root: TypeTreeNode = cls(-1, "", "", 0, 0, 0, [])
122+
fake_root: TypeTreeNode = cls(-1, "", "", 0, 0, [])
123123
stack: List[TypeTreeNode] = [fake_root]
124124
parent = fake_root
125125
prev = fake_root
126126

127-
for dict_node in nodes:
128-
node = cls(**dict_node)
127+
for node in nodes:
128+
if isinstance(node, dict):
129+
node = cls(node)
129130

130131
if node.m_Level > prev.m_Level:
131132
stack.append(parent)

‎UnityPyBoost/TypeTreeHelper.cpp‎

Lines changed: 18 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -586,23 +586,23 @@ static int TypeTreeNode_init(TypeTreeNodeObject *self, PyObject *args, PyObject
586586
"m_Type",
587587
"m_Name",
588588
"m_ByteSize",
589-
"m_TypeFlags",
590589
"m_Version",
591590
"m_Children",
591+
"m_TypeFlags",
592592
"m_VariableCount",
593593
"m_Index",
594594
"m_MetaFlag",
595595
"m_RefTypeHash",
596596
NULL};
597597

598-
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OOOOOO|OOOOO", (char **)kwlist,
598+
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OOOOO|OOOOOO", (char **)kwlist,
599599
&self->m_Level,
600600
&self->m_Type,
601601
&self->m_Name,
602602
&self->m_ByteSize,
603-
&self->m_TypeFlags,
604603
&self->m_Version,
605604
&self->m_Children,
605+
&self->m_TypeFlags,
606606
&self->m_VariableCount,
607607
&self->m_Index,
608608
&self->m_MetaFlag,
@@ -615,7 +615,6 @@ static int TypeTreeNode_init(TypeTreeNodeObject *self, PyObject *args, PyObject
615615
Py_INCREF(self->m_Type);
616616
Py_INCREF(self->m_Name);
617617
Py_INCREF(self->m_ByteSize);
618-
Py_INCREF(self->m_TypeFlags);
619618
Py_INCREF(self->m_Version);
620619
// optional fields
621620
if (self->m_Children == nullptr)
@@ -626,39 +625,23 @@ static int TypeTreeNode_init(TypeTreeNodeObject *self, PyObject *args, PyObject
626625
{
627626
Py_XINCREF(self->m_Children);
628627
}
629-
if (self->m_VariableCount == nullptr)
630-
{
631-
self->m_VariableCount = Py_None;
632-
}
633-
else
634-
{
635-
Py_INCREF(self->m_VariableCount);
636-
}
637-
if (self->m_Index == nullptr)
638-
{
639-
self->m_Index = Py_None;
640-
}
641-
else
642-
{
643-
Py_INCREF(self->m_Index);
644-
}
645-
if (self->m_MetaFlag == nullptr)
646-
{
647-
self->m_MetaFlag = Py_None;
648-
}
649-
else
650-
{
651-
Py_INCREF(self->m_MetaFlag);
652-
}
653-
if (self->m_RefTypeHash == nullptr)
654-
{
655-
self->m_RefTypeHash = Py_None;
656-
}
657-
else
658-
{
659-
Py_INCREF(self->m_RefTypeHash);
628+
629+
#define SET_NONE_IF_NULL(field) \
630+
if (self->field == nullptr) \
631+
{ \
632+
self->field = Py_None; \
633+
} \
634+
else \
635+
{ \
636+
Py_INCREF(self->field); \
660637
}
661638

639+
SET_NONE_IF_NULL(m_TypeFlags);
640+
SET_NONE_IF_NULL(m_VariableCount);
641+
SET_NONE_IF_NULL(m_Index);
642+
SET_NONE_IF_NULL(m_MetaFlag);
643+
SET_NONE_IF_NULL(m_RefTypeHash);
644+
662645
if (self->m_MetaFlag != Py_None && PyLong_AsLong(self->m_MetaFlag) & 0x4000)
663646
{
664647
self->_align = true;

‎UnityPyBoost/TypeTreeHelper.hpp‎

Lines changed: 1 addition & 1 deletion

0 commit comments

Comments
 (0)