From d936942d28972bf4288f034a382d7767fc7720aa Mon Sep 17 00:00:00 2001
From: hakril <project@hakril.net>
Date: Sat, 24 Nov 2018 22:46:58 +0100
Subject: [PATCH] [WIP] improving Token object possibilities

---
 tests/test_improved_buffer.py | 38 +++++++++++++++++++++++++++++++++++
 windows/winobject/token.py    |  2 +-
 2 files changed, 39 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_improved_buffer.py

diff --git a/tests/test_improved_buffer.py b/tests/test_improved_buffer.py
new file mode 100644
index 0000000..5829af9
--- /dev/null
+++ b/tests/test_improved_buffer.py
@@ -0,0 +1,38 @@
+import ctypes
+
+import windows.utils as utils
+import windows.generated_def as gdef
+
+from pfwtest import *
+
+
+@pytest.mark.parametrize("type, size", [
+(gdef.BYTE, 10),
+(gdef.DWORD, 12),
+(gdef.GUID, 42)])
+def test_improved_BUFFER_size(type, size):
+    x = utils.BUFFER(type, 1)(size=size)
+    assert len(x) == 1
+    assert x.real_size == size
+    assert len(x._raw_buffer_) == size
+
+@pytest.mark.parametrize("params, expected_type, expected_size", [
+(([gdef.DWORD(x) for x in range(10)],), gdef.DWORD, 10), # list
+(([gdef.GUID() for x in range(10)],), gdef.GUID, 10), # generator
+((range(42), gdef.ULONGLONG), gdef.ULONGLONG, 42), # Explicite type
+])
+def test_improved_buffer(params, expected_type, expected_size):
+    x = utils.buffer(*params)
+    assert x._type_ == expected_type
+    assert len(x) == expected_size
+
+@pytest.mark.parametrize("c_type, buffer, expected_size", [
+(gdef.CHAR, "12345", 5),
+(gdef.WCHAR, "\x001\x002\x003\x004\x005", 5),
+(gdef.DWORD, "1111222233334444", 4),
+])
+def test_parial_buffer_size_guess(c_type, buffer, expected_size):
+    buf = windows.utils.BUFFER(c_type).from_buffer_copy(buffer)
+    assert len(buf) == expected_size
+
+
diff --git a/windows/winobject/token.py b/windows/winobject/token.py
index 4c19289..d34e16e 100644
--- a/windows/winobject/token.py
+++ b/windows/winobject/token.py
@@ -122,7 +122,7 @@ class Token(utils.AutoHandle):
     def _get_token_infomations(self, infos_class, rtype):
         required_size = self._get_required_token_information_size(infos_class)
         requested_size = max(required_size, ctypes.sizeof(rtype))
-        buffer = utils.buffer_type(rtype, 1)(size=requested_size)
+        buffer = utils.BUFFER(rtype, 1)(size=requested_size)
         cbsize = gdef.DWORD()
         winproxy.GetTokenInformation(self.handle, infos_class, buffer, buffer.real_size, cbsize)
         return buffer[0]