cpp implementation for ArchiveStorageDecryptor.decrypt_block

UnityPy/UnityPyBoost.pyi

@@ -37,3 +37,10 @@ class TypeTreeNode:
     m_MetaFlag: Optional[int] = None
     m_RefTypeHash: Optional[int] = None
     _clean_name: str
+
+def decrypt_block(
+    index_bytes: bytes,
+    substitute_bytes: bytes,
+    data: Union[bytes, bytearray],
+    index: int,
+) -> None: ...

UnityPy/helpers/ArchiveStorageManager.py

@@ -4,6 +4,11 @@
 
 from ..streams import EndianBinaryReader
 
+try:
+    from UnityPy import UnityPyBoost
+except ImportError:
+    UnityPyBoost = None
+
 UNITY3D_SIGNATURE = b"#$unity3dchina!@"
 DECRYPT_KEY: bytes = None
 
@@ -101,6 +106,11 @@ def __init__(self, reader: EndianBinaryReader) -> None:
         )
 
     def decrypt_block(self, data: bytes, index: int):
+
+        if UnityPyBoost:
+            UnityPyBoost.decrypt_block(bytes(self.index), bytes(self.substitute), data, index)
+            return data
+
         offset = 0
         size = len(data)
         data = bytearray(data)

UnityPyBoost/ArchiveStorageDecryptor.cpp

@@ -0,0 +1,89 @@
+// based on https://github.com/RazTools/Studio/blob/main/AssetStudio/Crypto/UnityCN.cs
+
+#include "ArchiveStorageDecryptor.hpp"
+#include <Python.h>
+
+inline unsigned char decrypt_byte(unsigned char* bytes, uint64_t& offset, uint64_t& index, unsigned char* index_data, unsigned char* substitute_data)
+{
+    unsigned char count_byte = substitute_data[((index >> 2) & 3) + 4]
+                            + substitute_data[index & 3]
+                            + substitute_data[((index >> 4) & 3) + 8]
+                            + substitute_data[((unsigned char)index >> 6) + 12];
+    bytes[offset] = (unsigned char)((index_data[bytes[offset] & 0xF] - count_byte) & 0xF | 0x10 * (index_data[bytes[offset] >> 4] - count_byte));
+    count_byte = bytes[offset++];
+    index++;
+    return count_byte;
+}
+
+inline uint64_t decrypt(unsigned char* bytes, uint64_t index, uint64_t remaining, unsigned char* index_data, unsigned char* substitute_data)
+{
+    uint64_t offset = 0;
+
+    unsigned char current_byte = decrypt_byte(bytes, offset, index, index_data, substitute_data);
+    uint64_t current_byte_high = current_byte >> 4;
+    uint64_t current_byte_low = current_byte & 0xF;
+
+    if (current_byte_high == 0xF)
+    {
+        unsigned char count_byte;
+        do
+        {
+            count_byte = decrypt_byte(bytes, offset, index, index_data, substitute_data);
+            current_byte_high += count_byte;
+        } while (count_byte == 0xFF);
+    }
+
+    offset += current_byte_high;
+
+    if (offset < remaining)
+    {
+        decrypt_byte(bytes, offset, index, index_data, substitute_data);
+        decrypt_byte(bytes, offset, index, index_data, substitute_data);
+        if (current_byte_low == 0xF)
+        {
+            unsigned char count_byte;
+            do
+            {
+                count_byte = decrypt_byte(bytes, offset, index, index_data, substitute_data);
+            } while (count_byte == 0xFF);
+        }
+    }
+
+    return offset;
+}
+
+PyObject* decrypt_block(PyObject* self, PyObject* args) {
+    PyObject* py_index_bytes;
+    PyObject* py_substitute_bytes;
+    PyObject* py_data;
+    uint64_t index;
+
+    if (!PyArg_ParseTuple(args, "OOOi", &py_index_bytes, &py_substitute_bytes, &py_data, &index)) {
+        return NULL;
+    }
+
+    Py_buffer view;
+    if (PyObject_GetBuffer(py_data, &view, PyBUF_SIMPLE) != 0) {
+        return NULL;
+    }
+
+    if (!PyBytes_Check(py_index_bytes) || !PyBytes_Check(py_substitute_bytes)) {
+        PyBuffer_Release(&view);
+        PyErr_SetString(PyExc_TypeError, "Attributes 'index' and 'substitute' must be bytes");
+        return NULL;
+    }
+
+    unsigned char* data = (unsigned char*)view.buf;
+    uint64_t size = (uint64_t)view.len;
+    unsigned char* index_data = (unsigned char*)PyBytes_AS_STRING(py_index_bytes);
+    unsigned char* substitute_data = (unsigned char*)PyBytes_AS_STRING(py_substitute_bytes);
+
+    uint64_t offset = 0;
+    while (offset < size) {
+        offset += decrypt(data + offset, index++, size - offset, index_data, substitute_data);
+    }
+
+    PyBuffer_Release(&view);
+    Py_RETURN_NONE;
+}
+

UnityPyBoost/ArchiveStorageDecryptor.hpp

@@ -0,0 +1,5 @@
+#define PY_SSIZE_T_CLEAN
+#pragma once
+#include <Python.h>
+
+PyObject* decrypt_block(PyObject* self, PyObject* args);

UnityPyBoost/UnityPyBoost.cpp

@@ -2,6 +2,7 @@
 #include <Python.h>
 #include "Mesh.hpp"
 #include "TypeTreeHelper.hpp"
+#include "ArchiveStorageDecryptor.hpp"
 
 /* Mesh.py */
 
@@ -14,6 +15,10 @@ static struct PyMethodDef method_table[] = {
      (PyCFunction)read_typetree,
      METH_VARARGS | METH_KEYWORDS,
      "replacement for TypeTreeHelper.read_typetree"},
+     {"decrypt_block",
+     (PyCFunction)decrypt_block,
+     METH_VARARGS,
+     "replacement for ArchiveStorageDecryptor.decrypt_block"},
     {NULL,
      NULL,
      0,
@@ -39,4 +44,4 @@ PyMODINIT_FUNC PyInit_UnityPyBoost(void)
     PyObject *module = PyModule_Create(&UnityPyBoost_module);
     add_typetreenode_to_module(module);
     return module;
-}
+}