Coverage for agentos/tools/key_rotation.py: 0%

1""" 

2KeyRotation — automated secret key rotation with grace periods and scheduled callbacks. 

3 

4Supports: 

5 - Schedule-based rotation (interval in seconds) 

6 - Manual rotation trigger 

7 - Grace period (old key still valid for verification) 

8 - Current / pending / expired key states 

9 - Rotation hooks (pre_rotate, post_rotate) 

10 - Thread-safe 

11""" 

12 

13from __future__ import annotations 

14 

15import secrets 

16import threading 

17import time 

18from collections.abc import Callable 

19from dataclasses import dataclass 

20from enum import Enum 

21from typing import Any 

22 

23# ============================================================================ 

24# Key State 

25# ============================================================================ 

26 

27 

28class KeyState(Enum): 

29 CURRENT = "current" 

30 PENDING = "pending" # newly rotated, still in grace 

31 EXPIRED = "expired" 

32 

33 

34@dataclass 

35class KeyEntry: 

36 key: str 

37 state: KeyState 

38 created_at: float 

39 expires_at: float | None = None 

40 

41 

42# ============================================================================ 

43# KeyRotation 

44# ============================================================================ 

45 

46 

47class KeyRotation: 

48 """Automated secret key rotation with grace periods. 

49 

50 Usage: 

51 kr = KeyRotation(rotation_interval=3600, grace_period=300, key_length=32) 

52 kr.start() 

53 

54 current = kr.current_key # Use this for signing/encryption 

55 active_keys = kr.active_keys # All currently valid keys 

56 

57 # When a new key is rotated in: 

58 # - current becomes the new key 

59 # - old key stays in active_keys during grace period (for validation) 

60 # - after grace period expires, old key is removed 

61 """ 

62 

63 def __init__( 

64 self, 

65 rotation_interval: float = 3600.0, 

66 grace_period: float = 300.0, 

67 key_length: int = 32, 

68 ): 

69 if rotation_interval <= 0: 

70 raise ValueError("rotation_interval must be positive") 

71 if grace_period < 0: 

72 raise ValueError("grace_period must be non-negative") 

73 self._interval = rotation_interval 

74 self._grace_period = grace_period 

75 self._key_length = key_length 

76 self._keys: list[KeyEntry] = [] 

77 self._lock = threading.RLock() 

78 self._timer: threading.Timer | None = None 

79 self._running = False 

80 

81 # Hooks 

82 self._pre_rotate: list[Callable[[], None]] = [] 

83 self._post_rotate: list[Callable[[str, str], None]] = [] # (old_key, new_key) 

84 

85 # Seed with initial key 

86 self._rotate_now() 

87 

88 # ---------- Lifecycle ---------- 

89 

90 def start(self) -> None: 

91 with self._lock: 

92 if self._running: 

93 return 

94 self._running = True 

95 self._schedule_next() 

96 

97 def stop(self) -> None: 

98 with self._lock: 

99 self._running = False 

100 if self._timer: 

101 self._timer.cancel() 

102 self._timer = None 

103 

104 def _schedule_next(self) -> None: 

105 with self._lock: 

106 if not self._running: 

107 return 

108 self._timer = threading.Timer(self._interval, self._on_timer) 

109 self._timer.daemon = True 

110 self._timer.start() 

111 

112 def _on_timer(self) -> None: 

113 self._rotate_now() 

114 self._schedule_next() 

115 

116 # ---------- Rotation ---------- 

117 

118 def rotate(self) -> str: 

119 """Manually trigger a rotation. Returns the new key.""" 

120 return self._rotate_now() 

121 

122 def _rotate_now(self) -> str: 

123 new_key = secrets.token_hex(self._key_length) 

124 now = time.time() 

125 

126 with self._lock: 

127 old_key = self._keys[0].key if self._keys else None 

128 

129 # Notify pre-rotation 

130 self._notify_pre_rotate() 

131 

132 # Move current → pending (if grace > 0), otherwise expired 

133 for entry in self._keys: 

134 if entry.state == KeyState.CURRENT: 

135 if self._grace_period > 0: 

136 entry.state = KeyState.PENDING 

137 entry.expires_at = now + self._grace_period 

138 else: 

139 entry.state = KeyState.EXPIRED 

140 

141 # Add new current key 

142 new_entry = KeyEntry(key=new_key, state=KeyState.CURRENT, created_at=now) 

143 self._keys.insert(0, new_entry) 

144 

145 # Clean up expired 

146 self._keys = [e for e in self._keys if e.state != KeyState.EXPIRED] 

147 

148 # Notify post-rotation 

149 self._notify_post_rotate(old_key, new_key) 

150 

151 return new_key 

152 

153 # ---------- Key Access ---------- 

154 

155 @property 

156 def current_key(self) -> str | None: 

157 with self._lock: 

158 for entry in self._keys: 

159 if entry.state == KeyState.CURRENT: 

160 return entry.key 

161 return None 

162 

163 @property 

164 def active_keys(self) -> list[str]: 

165 """All currently valid keys (CURRENT + PENDING).""" 

166 with self._lock: 

167 self._cleanup_expired() 

168 return [e.key for e in self._keys if e.state in (KeyState.CURRENT, KeyState.PENDING)] 

169 

170 @property 

171 def pending_keys(self) -> list[str]: 

172 """Keys in grace period only.""" 

173 with self._lock: 

174 self._cleanup_expired() 

175 return [e.key for e in self._keys if e.state == KeyState.PENDING] 

176 

177 def is_valid(self, key: str) -> bool: 

178 """Check if a key is currently valid (CURRENT or PENDING).""" 

179 return key in self.active_keys 

180 

181 def _cleanup_expired(self) -> None: 

182 now = time.time() 

183 self._keys = [ 

184 e 

185 for e in self._keys 

186 if e.state != KeyState.EXPIRED and (e.expires_at is None or e.expires_at > now) 

187 ] 

188 

189 # ---------- Hooks ---------- 

190 

191 def on_pre_rotate(self, callback: Callable[[], None]) -> None: 

192 self._pre_rotate.append(callback) 

193 

194 def on_post_rotate(self, callback: Callable[[str, str], None]) -> None: 

195 self._post_rotate.append(callback) 

196 

197 def _notify_pre_rotate(self) -> None: 

198 for cb in self._pre_rotate: 

199 try: 

200 cb() 

201 except Exception: 

202 pass 

203 

204 def _notify_post_rotate(self, old_key: str | None, new_key: str) -> None: 

205 for cb in self._post_rotate: 

206 try: 

207 cb(old_key, new_key) 

208 except Exception: 

209 pass 

210 

211 # ---------- Info ---------- 

212 

213 @property 

214 def stats(self) -> dict[str, Any]: 

215 with self._lock: 

216 self._cleanup_expired() 

217 return { 

218 "total_keys": len(self._keys), 

219 "current": 1 if any(e.state == KeyState.CURRENT for e in self._keys) else 0, 

220 "pending": sum(1 for e in self._keys if e.state == KeyState.PENDING), 

221 "rotation_interval": self._interval, 

222 "grace_period": self._grace_period, 

223 "key_length": self._key_length, 

224 "running": self._running, 

225 }