some fix for Python api usage according to comment

K0lb3 · Oct 27, 2024 · 65deb2a · 65deb2a
1 parent a3a5cd5
commit 65deb2a
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 27 deletions.
diff --git a/UnityPy/UnityPyBoost.pyi b/UnityPy/UnityPyBoost.pyi
@@ -39,8 +39,8 @@ class TypeTreeNode:
     _clean_name: str
 
 def decrypt_block(
-    index_bytes: bytes,
-    substitute_bytes: bytes,
+    index_bytes: Union[bytes, bytearray],
+    substitute_bytes: Union[bytes, bytearray],
     data: Union[bytes, bytearray],
     index: int,
-) -> bytearray: ...
+) -> bytes: ...
diff --git a/UnityPy/helpers/ArchiveStorageManager.py b/UnityPy/helpers/ArchiveStorageManager.py
@@ -108,7 +108,7 @@ def __init__(self, reader: EndianBinaryReader) -> None:
     def decrypt_block(self, data: bytes, index: int):
 
         if UnityPyBoost:
-            return UnityPyBoost.decrypt_block(bytes(self.index), bytes(self.substitute), data, index)
+            return UnityPyBoost.decrypt_block(self.index, self.substitute, data, index)
 
         offset = 0
         size = len(data)

diff --git a/UnityPyBoost/ArchiveStorageDecryptor.cpp b/UnityPyBoost/ArchiveStorageDecryptor.cpp
@@ -3,7 +3,7 @@
 #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)
+inline unsigned char decrypt_byte(unsigned char *bytes, uint64_t& offset, uint64_t& index, const unsigned char *index_data, const unsigned char *substitute_data)
 {
     unsigned char count_byte = substitute_data[((index >> 2) & 3) + 4]
                             + substitute_data[index & 3]
@@ -15,7 +15,7 @@ inline unsigned char decrypt_byte(unsigned char *bytes, uint64_t& offset, uint64
     return count_byte;
 }
 
-inline uint64_t decrypt(unsigned char *bytes, uint64_t index, uint64_t remaining, unsigned char *index_data, unsigned char *substitute_data)
+inline uint64_t decrypt(unsigned char *bytes, uint64_t index, uint64_t remaining, const unsigned char *index_data, const unsigned char *substitute_data)
 {
     uint64_t offset = 0;
 
@@ -53,39 +53,37 @@ inline uint64_t decrypt(unsigned char *bytes, uint64_t index, uint64_t remaining
 }
 
 PyObject *decrypt_block(PyObject *self, PyObject *args) {
-    PyObject *py_index_bytes;
-    PyObject *py_substitute_bytes;
-    PyObject *py_data;
+    Py_buffer index_data;
+    Py_buffer substitute_data;
+    Py_buffer data;
     uint64_t index;
 
-    if (!PyArg_ParseTuple(args, "OOOi", &py_index_bytes, &py_substitute_bytes, &py_data, &index)) {
+    if (!PyArg_ParseTuple(args, "y*y*y*K", &index_data, &substitute_data, &data, &index)) {
+        if (index_data.buf) PyBuffer_Release(&index_data);
+        if (substitute_data.buf) PyBuffer_Release(&substitute_data);
+        if (data.buf) PyBuffer_Release(&data);
         return NULL;
     }
 
-    PyObject *result = PyBytes_FromObject(py_data);
-
-    Py_buffer view;
-    if (PyObject_GetBuffer(result, &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");
+    PyObject *result = PyBytes_FromStringAndSize(NULL, data.len);
+    if (result == NULL) {
+        PyBuffer_Release(&index_data);
+        PyBuffer_Release(&substitute_data);
+        PyBuffer_Release(&data);
         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);
+    unsigned char *result_raw = (unsigned char *)PyBytes_AS_STRING(result);
+    memcpy(result_raw, data.buf, data.len);
 
     uint64_t offset = 0;
-    while (offset < size) {
-        offset += decrypt(data + offset, index++, size - offset, index_data, substitute_data);
+    while (offset < data.len) {
+        offset += decrypt(result_raw + offset, index++, data.len - offset, (unsigned char *)index_data.buf, (unsigned char *)substitute_data.buf);
     }
 
-    PyBuffer_Release(&view);
+    PyBuffer_Release(&index_data);
+    PyBuffer_Release(&substitute_data);
+    PyBuffer_Release(&data);
 
     return result;
 }