Initial commit

quantumpipe/securebuf.py

@@ -0,0 +1,100 @@
+"""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()