Skip to content

Commit b1a4d21

Browse files
radumariasradumarias
committed
tests
1 parent 0f3cd09 commit b1a4d21

File tree

3 files changed

+110
-91
lines changed

3 files changed

+110
-91
lines changed

‎examples/lock_and_zeroize.py

Lines changed: 32 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -30,34 +30,35 @@ def unlock_memory(buffer):
3030
if MUNLOCK(address, size) != 0:
3131
raise RuntimeError("Failed to unlock memory")
3232

33-
34-
try:
35-
print("allocate memory")
36-
37-
# regular array
38-
arr = bytearray(b"1234567890")
39-
40-
# numpy array
41-
arr_np = np.array([0] * 10, dtype=np.uint8)
42-
arr_np[:] = arr
43-
assert arr_np.tobytes() == b"1234567890"
44-
45-
print("locking memory")
46-
47-
lock_memory(arr)
48-
lock_memory(arr_np)
49-
50-
print("zeroize'ing...: ")
51-
zeroize1(arr)
52-
zeroize_np(arr_np)
53-
54-
print("checking if is zeroized")
55-
assert arr == bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
56-
assert all(arr_np == 0)
57-
58-
print("all good, bye!")
59-
finally:
60-
# Unlock the memory
61-
print("unlocking memory")
62-
unlock_memory(arr)
63-
unlock_memory(arr_np)
33+
if __name__ == "__main__":
34+
try:
35+
print("allocate memory")
36+
37+
# regular array
38+
arr = bytearray(b"1234567890")
39+
40+
# numpy array
41+
arr_np = np.array([0] * 10, dtype=np.uint8)
42+
arr_np[:] = arr
43+
assert arr_np.tobytes() == b"1234567890"
44+
45+
print("locking memory")
46+
47+
lock_memory(arr)
48+
lock_memory(arr_np)
49+
50+
print("zeroize'ing...: ")
51+
zeroize1(arr)
52+
zeroize_np(arr_np)
53+
54+
print("checking if is zeroized")
55+
assert arr == bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
56+
assert all(arr_np == 0)
57+
58+
print("all good, bye!")
59+
60+
finally:
61+
# Unlock the memory
62+
print("unlocking memory")
63+
unlock_memory(arr)
64+
unlock_memory(arr_np)

‎examples/zeroize_before_fork.py

Lines changed: 23 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -31,24 +31,26 @@ def unlock_memory(buffer):
3131
raise RuntimeError("Failed to unlock memory")
3232

3333

34-
try:
35-
sensitive_data = bytearray(b"Sensitive Information")
36-
lock_memory(sensitive_data)
37-
38-
print("Before zeroization:", sensitive_data)
39-
40-
zeroize1(sensitive_data)
41-
print("After zeroization:", sensitive_data)
42-
43-
# Forking after zeroization to ensure no sensitive data is copied
44-
pid = os.fork()
45-
if pid == 0:
46-
# This is the child process
47-
print("Child process memory after fork:", sensitive_data)
48-
else:
49-
# This is the parent process
50-
os.wait() # Wait for the child process to exit
51-
finally:
52-
# Unlock the memory
53-
print("unlocking memory")
54-
unlock_memory(sensitive_data)
34+
if __name__ == "__main__":
35+
try:
36+
sensitive_data = bytearray(b"Sensitive Information")
37+
lock_memory(sensitive_data)
38+
39+
print("Before zeroization:", sensitive_data)
40+
41+
zeroize1(sensitive_data)
42+
print("After zeroization:", sensitive_data)
43+
44+
# Forking after zeroization to ensure no sensitive data is copied
45+
pid = os.fork()
46+
if pid == 0:
47+
# This is the child process
48+
print("Child process memory after fork:", sensitive_data)
49+
else:
50+
# This is the parent process
51+
os.wait() # Wait for the child process to exit
52+
53+
finally:
54+
# Unlock the memory
55+
print("unlocking memory")
56+
unlock_memory(sensitive_data)

‎tests/test_zeroize.py

Lines changed: 55 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,35 @@
1-
import ctypes
2-
import os
31
import unittest
42
import zeroize
53
import numpy as np
64

75

8-
# Lock memory using ctypes
9-
def lock_memory():
10-
libc = ctypes.CDLL("libc.so.6")
11-
# Lock all current and future pages from being swapped out
12-
libc.mlockall(ctypes.c_int(0x02 | 0x04)) # MCL_CURRENT | MCL_FUTURE
6+
import ctypes
7+
8+
9+
# Load the C standard library
10+
LIBC = ctypes.CDLL("libc.so.6")
11+
MLOCK = LIBC.mlock
12+
MUNLOCK = LIBC.munlock
13+
14+
# Define mlock and munlock argument types
15+
MLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
16+
MUNLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
17+
18+
19+
def lock_memory(buffer):
20+
"""Locks the memory of the given buffer."""
21+
address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
22+
size = len(buffer)
23+
if MLOCK(address, size) != 0:
24+
raise RuntimeError("Failed to lock memory")
1325

1426

15-
def unlock_memory():
16-
libc = ctypes.CDLL("libc.so.6")
17-
# Unlock all locked pages
18-
libc.munlockall()
27+
def unlock_memory(buffer):
28+
"""Unlocks the memory of the given buffer."""
29+
address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
30+
size = len(buffer)
31+
if MUNLOCK(address, size) != 0:
32+
raise RuntimeError("Failed to unlock memory")
1933

2034

2135
SIZES_MB = [
@@ -43,44 +57,46 @@ def unlock_memory():
4357
class TestStringMethods(unittest.TestCase):
4458

4559
def test_zeroize1(self):
46-
lock_memory()
47-
48-
arr = bytearray(b"1234567890")
49-
zeroize.zeroize1(arr)
50-
self.assertEqual(arr, bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"))
60+
try:
61+
arr = bytearray(b"1234567890")
62+
lock_memory(arr)
63+
zeroize.zeroize1(arr)
64+
self.assertEqual(
65+
arr, bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
66+
)
5167

52-
unlock_memory()
68+
finally:
69+
unlock_memory(arr)
5370

5471
def test_zeroize_np(self):
55-
lock_memory()
72+
try:
73+
arr = np.array([0] * 10, dtype=np.uint8)
74+
arr[:] = bytes(b"1234567890")
75+
zeroize.zeroize_np(arr)
76+
self.assertEqual(True, all(arr == 0))
5677

57-
arr = np.array([0] * 10, dtype=np.uint8)
58-
arr[:] = bytes(b"1234567890")
59-
zeroize.zeroize_np(arr)
60-
self.assertEqual(True, all(arr == 0))
61-
62-
unlock_memory()
78+
finally:
79+
unlock_memory(arr)
6380

6481
def test_zeroize1_sizes(self):
65-
# lock_memory()
66-
6782
for size in SIZES_MB:
68-
arr = bytearray(int(size * 1024 * 1024))
69-
zeroize.zeroize1(arr)
70-
self.assertEqual(arr, bytearray(int(size * 1024 * 1024)))
71-
72-
# unlock_memory()
83+
try:
84+
arr = bytearray(int(size * 1024 * 1024))
85+
zeroize.zeroize1(arr)
86+
self.assertEqual(arr, bytearray(int(size * 1024 * 1024)))
7387

74-
def test_zeroize_np_sizes(self):
75-
# lock_memory()
88+
finally:
89+
unlock_memory(arr)
7690

91+
def test_zeroize_np_sizes(self):
7792
for size in [size for size in SIZES_MB if size < 4]:
78-
array_size = int(size * 1024 * 1024)
79-
random_array = np.random.randint(0, 256, array_size, dtype=np.uint8)
80-
zeroize.zeroize_np(random_array)
81-
self.assertEqual(True, all(random_array == 0))
82-
83-
# unlock_memory()
93+
try:
94+
array_size = int(size * 1024 * 1024)
95+
random_array = np.random.randint(0, 256, array_size, dtype=np.uint8)
96+
zeroize.zeroize_np(random_array)
97+
self.assertEqual(True, all(random_array == 0))
98+
finally:
99+
unlock_memory(random_array)
84100

85101

86102
if __name__ == "__main__":

0 commit comments

Comments
 (0)