From a3a5cd549aef2e4940da55991c6ee026a3637202 Mon Sep 17 00:00:00 2001
From: AXiX-official <2879710747@qq.com>
Date: Sat, 26 Oct 2024 09:09:08 +0800
Subject: [PATCH] copy bytes by PyBytes_FromObject

---
 UnityPy/UnityPyBoost.pyi                 |  2 +-
 UnityPyBoost/ArchiveStorageDecryptor.cpp | 14 +++++---------
 2 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/UnityPy/UnityPyBoost.pyi b/UnityPy/UnityPyBoost.pyi
index ffbb20f2..8d829710 100644
--- a/UnityPy/UnityPyBoost.pyi
+++ b/UnityPy/UnityPyBoost.pyi
@@ -43,4 +43,4 @@ def decrypt_block(
     substitute_bytes: bytes,
     data: Union[bytes, bytearray],
     index: int,
-) -> bytes: ...
+) -> bytearray: ...
diff --git a/UnityPyBoost/ArchiveStorageDecryptor.cpp b/UnityPyBoost/ArchiveStorageDecryptor.cpp
index e26f58cc..fabe4d17 100644
--- a/UnityPyBoost/ArchiveStorageDecryptor.cpp
+++ b/UnityPyBoost/ArchiveStorageDecryptor.cpp
@@ -62,8 +62,10 @@ PyObject *decrypt_block(PyObject *self, PyObject *args) {
         return NULL;
     }
 
+    PyObject *result = PyBytes_FromObject(py_data);
+
     Py_buffer view;
-    if (PyObject_GetBuffer(py_data, &view, PyBUF_SIMPLE) != 0) {
+    if (PyObject_GetBuffer(result, &view, PyBUF_SIMPLE) != 0) {
         return NULL;
     }
 
@@ -73,22 +75,16 @@ PyObject *decrypt_block(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    const unsigned char *data = (unsigned char *)view.buf;
+    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 *decrypted_data = (unsigned char *)PyMem_Malloc(size + 1);
-    decrypted_data[size] = 0;
-    memcpy(decrypted_data, data, size);
-
     uint64_t offset = 0;
     while (offset < size) {
-        offset += decrypt(decrypted_data + offset, index++, size - offset, index_data, substitute_data);
+        offset += decrypt(data + offset, index++, size - offset, index_data, substitute_data);
     }
 
-    PyObject* result = PyBytes_FromStringAndSize((const char*)decrypted_data, size);
-    PyMem_Free(decrypted_data);
     PyBuffer_Release(&view);
 
     return result;