101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
"""SecureBuffer -- mlock'd, core-dump-excluded, wipe-on-close memory.
|
|
Unused by the package itself; kept for applications."""
|
|
|
|
import ctypes
|
|
import mmap
|
|
import os
|
|
from contextlib import ExitStack, suppress
|
|
|
|
|
|
def wipe(buf):
|
|
n = memoryview(buf).cast("B")
|
|
n[:] = b"\xff" * len(n)
|
|
|
|
|
|
class SecureBuffer:
|
|
def __init__(self, size: int):
|
|
with ExitStack() as defer:
|
|
self._mmap = mmap.mmap(-1, size, access=mmap.ACCESS_WRITE)
|
|
defer.callback(self._mmap.close)
|
|
self._lock()
|
|
defer.callback(self._unlock)
|
|
self._dontdump()
|
|
defer.pop_all()
|
|
|
|
@property
|
|
def closed(self):
|
|
return self._mmap.closed
|
|
|
|
def _addr(self):
|
|
return ctypes.addressof(ctypes.c_char.from_buffer(self._mmap))
|
|
|
|
def _lock(self):
|
|
addr = self._addr()
|
|
size = len(self._mmap)
|
|
if os.name == "nt":
|
|
api = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
api.VirtualLock.argtypes = (ctypes.c_void_p, ctypes.c_size_t)
|
|
api.VirtualLock.restype = ctypes.c_bool
|
|
if not api.VirtualLock(addr, size):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
else:
|
|
libc = ctypes.CDLL(None, use_errno=True)
|
|
libc.mlock.argtypes = (ctypes.c_void_p, ctypes.c_size_t)
|
|
libc.mlock.restype = ctypes.c_int
|
|
if libc.mlock(addr, size):
|
|
err = ctypes.get_errno()
|
|
raise OSError(err, os.strerror(err))
|
|
|
|
def _unlock(self):
|
|
addr = self._addr()
|
|
size = len(self._mmap)
|
|
if os.name == "nt":
|
|
api = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
api.VirtualUnlock.argtypes = (ctypes.c_void_p, ctypes.c_size_t)
|
|
api.VirtualUnlock.restype = ctypes.c_bool
|
|
if not api.VirtualUnlock(addr, size):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
else:
|
|
libc = ctypes.CDLL(None, use_errno=True)
|
|
libc.munlock.argtypes = (ctypes.c_void_p, ctypes.c_size_t)
|
|
libc.munlock.restype = ctypes.c_int
|
|
if libc.munlock(addr, size):
|
|
err = ctypes.get_errno()
|
|
raise OSError(err, os.strerror(err))
|
|
|
|
def _dontdump(self):
|
|
if hasattr(self._mmap, "madvise") and hasattr(mmap, "MADV_DONTDUMP"):
|
|
self._mmap.madvise(mmap.MADV_DONTDUMP)
|
|
|
|
def close(self):
|
|
if self.closed:
|
|
return
|
|
self._mmap[:] = b"\xff" * len(self._mmap)
|
|
self._unlock()
|
|
self._mmap.close()
|
|
|
|
def __buffer__(self, flags):
|
|
return memoryview(self._mmap)
|
|
|
|
def __getitem__(self, key):
|
|
return memoryview(self._mmap)[key]
|
|
|
|
def __setitem__(self, key, value):
|
|
memoryview(self._mmap)[key] = value
|
|
|
|
def __bytes__(self):
|
|
raise TypeError("SecureBuffer cannot be converted to bytes")
|
|
|
|
def __len__(self):
|
|
return len(self._mmap)
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
self.close()
|
|
|
|
def __del__(self):
|
|
with suppress(Exception):
|
|
self.close()
|