security/memory.py


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38


import nacl.bindings
import ctypes
import platform

class SecureMemory:
    def __init__(self):
        self.protected_blocks = []
        
    def protect_memory(self, data: bytes) -> bytearray:
        """Protect memory block from being swapped to disk"""
        protected = bytearray(data)
        self.protected_blocks.append(protected)
        
        if platform.system() != 'Windows':
            try:
                # Lock memory to prevent swapping
                libc = ctypes.CDLL('libc.so.6')
                libc.mlock(
                    ctypes.c_void_p(protected.__array_interface__['data'][0]),
                    len(protected)
                )
            except:
                pass
                
        return protected
        
    def secure_wipe(self, data: bytearray):
        """Securely wipe memory block"""
        if data in self.protected_blocks:
            self.protected_blocks.remove(data)
            
        nacl.bindings.sodium_memzero(data)
        
    def wipe_all(self):
        """Wipe all protected memory blocks"""
        for block in self.protected_blocks[:]:
            self.secure_wipe(block)