Coverage for src/taipanstack/security/decorators.py: 100%

1""" 

2Security decorators for robust Python applications. 

3 

4Provides decorators for input validation, exception handling, 

5timeout control, and other security patterns. Compatible with 

6any Python framework (Flask, FastAPI, Django, etc.). 

7""" 

8 

9import functools 

10import inspect 

11import math 

12import signal 

13import sys 

14import threading 

15from collections.abc import Callable 

16from types import FrameType 

17from typing import ParamSpec, TypeVar 

18 

19from taipanstack.security.guards import SecurityError 

20 

21P = ParamSpec("P") 

22R = TypeVar("R") 

23T = TypeVar("T") 

24 

25 

26class OperationTimeoutError(Exception): 

27 """Raised when a function exceeds its timeout limit.""" 

28 

29 def __init__(self, seconds: float, func_name: str = "function") -> None: 

30 """Initialize OperationTimeoutError. 

31 

32 Args: 

33 seconds: The timeout that was exceeded. 

34 func_name: Name of the function that timed out. 

35 

36 """ 

37 self.seconds = seconds 

38 self.func_name = func_name 

39 super().__init__(f"{func_name} timed out after {seconds} seconds") 

40 

41 

42class ValidationError(Exception): 

43 """Raised when input validation fails.""" 

44 

45 def __init__( 

46 self, 

47 message: str, 

48 param_name: str | None = None, 

49 value: object = None, 

50 ) -> None: 

51 """Initialize ValidationError. 

52 

53 Args: 

54 message: Description of the validation failure. 

55 param_name: Name of the parameter that failed. 

56 value: The invalid value (sanitized). 

57 

58 """ 

59 self.param_name = param_name 

60 self.value = value 

61 super().__init__(message) 

62 

63 

64def validate_inputs( 

65 **validators: Callable[[object], object], 

66) -> Callable[[Callable[P, R]], Callable[P, R]]: 

67 """Decorator to validate function inputs. 

68 

69 Validates function arguments using provided validator functions. 

70 Validators should raise ValueError or ValidationError on invalid input. 

71 

72 Args: 

73 **validators: Mapping of parameter names to validator functions. 

74 

75 Returns: 

76 Decorated function with input validation. 

77 

78 Example: 

79 >>> from taipanstack.security.validators import validate_email, validate_port 

80 >>> @validate_inputs(email=validate_email, port=validate_port) 

81 ... def connect(email: str, port: int) -> None: 

82 ... pass 

83 >>> connect(email="invalid", port=8080) 

84 ValidationError: Invalid email format: invalid 

85 

86 """ 

87 

88 def decorator(func: Callable[P, R]) -> Callable[P, R]: 

89 sig = inspect.signature(func) 

90 

91 @functools.wraps(func) 

92 def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: 

93 # Bind and apply defaults on each call 

94 bound = sig.bind(*args, **kwargs) 

95 bound.apply_defaults() 

96 

97 # Validate each parameter that has a validator 

98 for param_name, validator in validators.items(): 

99 if param_name in bound.arguments: 

100 value = bound.arguments[param_name] 

101 try: 

102 # Call validator - it should raise on invalid input 

103 validated = validator(value) 

104 # Update to validated value if returned 

105 if validated is not None: 

106 bound.arguments[param_name] = validated 

107 except (ValueError, TypeError) as e: 

108 raise ValidationError( 

109 str(e), 

110 param_name=param_name, 

111 value=repr(value)[:100], 

112 ) from e 

113 

114 # Call original function with validated arguments 

115 return func(*bound.args, **bound.kwargs) 

116 

117 return wrapper 

118 

119 return decorator 

120 

121 

122def guard_exceptions( 

123 *, 

124 catch: tuple[type[Exception], ...] = (Exception,), 

125 reraise_as: type[Exception] | None = None, 

126 default: T | None = None, 

127 log_errors: bool = True, 

128) -> Callable[[Callable[P, R]], Callable[P, R | T | None]]: 

129 """Decorator to safely handle exceptions. 

130 

131 Catches exceptions and optionally re-raises as a different type 

132 or returns a default value. 

133 

134 Args: 

135 catch: Exception types to catch. 

136 reraise_as: Exception type to re-raise as (None = don't reraise). 

137 default: Default value to return if exception caught and not reraised. 

138 log_errors: Whether to log caught exceptions. 

139 

140 Returns: 

141 Decorated function with exception handling. 

142 

143 Example: 

144 >>> @guard_exceptions(catch=(IOError,), reraise_as=SecurityError) 

145 ... def read_file(path: str) -> str: 

146 ... return open(path).read() 

147 >>> read_file("/nonexistent") 

148 SecurityError: [guard_exceptions] ... 

149 

150 """ 

151 

152 def decorator(func: Callable[P, R]) -> Callable[P, R | T | None]: 

153 @functools.wraps(func) 

154 def wrapper(*args: P.args, **kwargs: P.kwargs) -> R | T | None: 

155 try: 

156 return func(*args, **kwargs) 

157 except catch as e: 

158 if log_errors: 

159 import logging 

160 

161 logging.getLogger("taipanstack.security").warning( 

162 "Exception caught in %s: %s", 

163 func.__name__, 

164 str(e), 

165 ) 

166 

167 if reraise_as is not None: 

168 if reraise_as == SecurityError: 

169 raise SecurityError( 

170 str(e), 

171 guard_name="guard_exceptions", 

172 ) from e 

173 raise reraise_as(str(e)) from e 

174 

175 return default 

176 

177 return wrapper 

178 

179 return decorator 

180 

181 

182def timeout( 

183 seconds: float, 

184 *, 

185 use_signal: bool = True, 

186) -> Callable[[Callable[P, R]], Callable[P, R]]: 

187 """Decorator to limit function execution time. 

188 

189 Uses signal-based timeout on Unix or thread-based on Windows. 

190 Signal-based is more reliable but only works in main thread. 

191 

192 Args: 

193 seconds: Maximum execution time in seconds. 

194 use_signal: Use signal-based timeout (Unix only, main thread only). 

195 

196 Returns: 

197 Decorated function with timeout. 

198 

199 Example: 

200 >>> @timeout(5.0) 

201 ... def slow_operation() -> str: 

202 ... import time 

203 ... time.sleep(10) 

204 ... return "done" 

205 >>> slow_operation() 

206 TimeoutError: slow_operation timed out after 5.0 seconds 

207 

208 """ 

209 # Security Enhancement: explicitly validate bounds using math.isfinite() 

210 # and check for non-negative limits to prevent silent NaN propagation, 

211 # unhandled ValueError exceptions from threading/asyncio primitives, 

212 # or unexpected infinite blocking behaviors. 

213 if ( 

214 not isinstance(seconds, (int, float)) 

215 or not math.isfinite(seconds) 

216 or seconds < 0 

217 ): 

218 raise ValueError("timeout must be a finite non-negative number") 

219 

220 def decorator(func: Callable[P, R]) -> Callable[P, R]: 

221 @functools.wraps(func) 

222 def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: 

223 # Determine if we can use signals 

224 can_use_signal = ( 

225 use_signal 

226 and sys.platform != "win32" 

227 and threading.current_thread() is threading.main_thread() 

228 ) 

229 

230 if can_use_signal: 

231 return _timeout_with_signal( 

232 func, 

233 seconds, 

234 args, 

235 dict(kwargs), 

236 ) 

237 return _timeout_with_thread( 

238 func, 

239 seconds, 

240 args, 

241 dict(kwargs), 

242 ) 

243 

244 return wrapper 

245 

246 return decorator 

247 

248 

249def _timeout_with_signal( 

250 func: Callable[..., R], 

251 seconds: float, 

252 args: tuple[object, ...], 

253 kwargs: dict[str, object], 

254) -> R: 

255 """Implement timeout using Unix signals.""" 

256 

257 def handler(_signum: int, _frame: FrameType | None) -> None: 

258 raise OperationTimeoutError(seconds, func.__name__) 

259 

260 # Set up signal handler 

261 old_handler = signal.signal(signal.SIGALRM, handler) 

262 signal.setitimer(signal.ITIMER_REAL, seconds) 

263 

264 try: 

265 return func(*args, **kwargs) 

266 finally: 

267 # Restore old handler and cancel alarm 

268 signal.setitimer(signal.ITIMER_REAL, 0) 

269 signal.signal(signal.SIGALRM, old_handler) 

270 

271 

272def _timeout_with_thread( 

273 func: Callable[..., R], 

274 seconds: float, 

275 args: tuple[object, ...], 

276 kwargs: dict[str, object], 

277) -> R: 

278 """Implement timeout using a separate thread.""" 

279 result: list[R] = [] 

280 exception: list[BaseException] = [] 

281 

282 def target() -> None: 

283 try: 

284 result.append(func(*args, **kwargs)) 

285 except BaseException as e: 

286 exception.append(e) 

287 

288 thread = threading.Thread(target=target) 

289 thread.daemon = True 

290 thread.start() 

291 thread.join(timeout=seconds) 

292 

293 if thread.is_alive(): 

294 # Thread still running - timeout occurred 

295 raise OperationTimeoutError(seconds, func.__name__) 

296 

297 if exception: 

298 raise exception[0] 

299 

300 return result[0] 

301 

302 

303def deprecated( 

304 message: str = "", 

305 *, 

306 removal_version: str | None = None, 

307) -> Callable[[Callable[P, R]], Callable[P, R]]: 

308 """Mark a function as deprecated. 

309 

310 Emits a warning when the decorated function is called. 

311 

312 Args: 

313 message: Additional deprecation message. 

314 removal_version: Version when function will be removed. 

315 

316 Returns: 

317 Decorated function that warns on use. 

318 

319 Example: 

320 >>> @deprecated("Use new_function instead", removal_version="2.0") 

321 ... def old_function() -> None: 

322 ... pass 

323 

324 """ 

325 

326 def decorator(func: Callable[P, R]) -> Callable[P, R]: 

327 @functools.wraps(func) 

328 def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: 

329 import warnings 

330 

331 msg = f"{func.__name__} is deprecated." 

332 if removal_version: 

333 msg += f" Will be removed in version {removal_version}." 

334 if message: 

335 msg += f" {message}" 

336 

337 warnings.warn(msg, DeprecationWarning, stacklevel=2) 

338 return func(*args, **kwargs) 

339 

340 return wrapper 

341 

342 return decorator 

343 

344 

345def require_type( 

346 **type_hints: type, 

347) -> Callable[[Callable[P, R]], Callable[P, R]]: 

348 """Decorator to enforce runtime type checking. 

349 

350 Validates that arguments match specified types at runtime. 

351 

352 Args: 

353 **type_hints: Mapping of parameter names to expected types. 

354 

355 Returns: 

356 Decorated function with type checking. 

357 

358 Example: 

359 >>> @require_type(name=str, count=int) 

360 ... def greet(name: str, count: int) -> None: 

361 ... print(f"Hello {name}" * count) 

362 >>> greet(name=123, count=2) 

363 TypeError: Parameter 'name' expected str, got int 

364 

365 """ 

366 

367 def decorator(func: Callable[P, R]) -> Callable[P, R]: 

368 sig = inspect.signature(func) 

369 

370 @functools.wraps(func) 

371 def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: 

372 bound = sig.bind(*args, **kwargs) 

373 bound.apply_defaults() 

374 

375 for param_name, expected_type in type_hints.items(): 

376 if param_name in bound.arguments: 

377 value = bound.arguments[param_name] 

378 if not isinstance(value, expected_type): 

379 raise TypeError( 

380 f"Parameter '{param_name}' expected " 

381 f"{expected_type.__name__}, got {type(value).__name__}", 

382 ) 

383 

384 return func(*bound.args, **bound.kwargs) 

385 

386 return wrapper 

387 

388 return decorator