Coverage for src / kdbxtool / security / memory.py: 100%

1"""Secure memory handling for sensitive data. 

2 

3This module provides SecureBytes, a mutable byte container that: 

4- Stores data in a bytearray (mutable, can be zeroized) 

5- Automatically zeroizes memory on destruction 

6- Supports context manager protocol for guaranteed cleanup 

7- Prevents accidental exposure through repr/str 

8""" 

9 

10from __future__ import annotations 

11 

12from typing import Self 

13 

14 

15class SecureBytes: 

16 """A secure container for sensitive byte data that zeroizes on destruction. 

17 

18 Unlike Python's immutable `bytes`, SecureBytes uses a mutable `bytearray` 

19 that can be explicitly zeroed when no longer needed. This prevents sensitive 

20 data like passwords and cryptographic keys from lingering in memory. 

21 

22 Usage: 

23 # Basic usage 

24 key = SecureBytes(derived_key_bytes) 

25 # ... use key.data for crypto operations ... 

26 key.zeroize() # Explicit cleanup 

27 

28 # Context manager (recommended) 

29 with SecureBytes(password.encode()) as pwd: 

30 hash = sha256(pwd.data) 

31 # Automatically zeroized here 

32 

33 Note: 

34 While this provides defense-in-depth against memory disclosure attacks, 

35 Python's memory management means copies may still exist. For maximum 

36 security, consider using specialized libraries like PyNaCl's SecretBox. 

37 """ 

38 

39 __slots__ = ("_buffer", "_zeroized") 

40 

41 def __init__(self, data: bytes | bytearray) -> None: 

42 """Initialize with sensitive data. 

43 

44 Args: 

45 data: The sensitive bytes to protect. Will be copied into internal buffer. 

46 """ 

47 self._buffer = bytearray(data) 

48 self._zeroized = False 

49 

50 @property 

51 def data(self) -> bytes: 

52 """Access the underlying data as immutable bytes. 

53 

54 Returns: 

55 The protected data as bytes. 

56 

57 Raises: 

58 ValueError: If the buffer has already been zeroized. 

59 """ 

60 if self._zeroized: 

61 raise ValueError("SecureBytes has been zeroized") 

62 return bytes(self._buffer) 

63 

64 def __len__(self) -> int: 

65 """Return the length of the protected data.""" 

66 return len(self._buffer) 

67 

68 def __bool__(self) -> bool: 

69 """Return True if buffer contains data and hasn't been zeroized.""" 

70 return not self._zeroized and len(self._buffer) > 0 

71 

72 def zeroize(self) -> None: 

73 """Overwrite the buffer with zeros. 

74 

75 This method overwrites every byte in the buffer with 0x00, 

76 making the original data unrecoverable from this object. 

77 Safe to call multiple times. 

78 """ 

79 if not self._zeroized: 

80 for i in range(len(self._buffer)): 

81 self._buffer[i] = 0 

82 self._zeroized = True 

83 

84 def __del__(self) -> None: 

85 """Ensure buffer is zeroized when object is garbage collected.""" 

86 self.zeroize() 

87 

88 def __enter__(self) -> Self: 

89 """Context manager entry.""" 

90 return self 

91 

92 def __exit__( 

93 self, 

94 exc_type: type[BaseException] | None, 

95 exc_val: BaseException | None, 

96 exc_tb: object, 

97 ) -> None: 

98 """Context manager exit - always zeroize.""" 

99 self.zeroize() 

100 

101 def __repr__(self) -> str: 

102 """Safe repr that doesn't expose data.""" 

103 if self._zeroized: 

104 return "SecureBytes(<zeroized>)" 

105 return f"SecureBytes(<{len(self._buffer)} bytes>)" 

106 

107 def __str__(self) -> str: 

108 """Safe str that doesn't expose data.""" 

109 return self.__repr__() 

110 

111 def __eq__(self, other: object) -> bool: 

112 """Constant-time comparison to prevent timing attacks. 

113 

114 Uses hmac.compare_digest for constant-time comparison. 

115 """ 

116 if not isinstance(other, SecureBytes): 

117 return NotImplemented 

118 if self._zeroized or other._zeroized: 

119 return False 

120 

121 import hmac 

122 

123 return hmac.compare_digest(self._buffer, other._buffer) 

124 

125 def __hash__(self) -> int: 

126 """Raise TypeError - SecureBytes should not be hashable.""" 

127 raise TypeError("SecureBytes is not hashable") 

128 

129 @classmethod 

130 def from_str(cls, s: str, encoding: str = "utf-8") -> Self: 

131 """Create SecureBytes from a string. 

132 

133 Args: 

134 s: String to encode 

135 encoding: Character encoding (default: utf-8) 

136 

137 Returns: 

138 SecureBytes containing the encoded string 

139 """ 

140 return cls(s.encode(encoding))