"""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()