Coverage for agentos/enterprise/api_keys.py: 49%

1""" 

2AgentOS Enterprise — API Key Management. 

3 

4功能： 

5 - API Key 创建/撤销/轮转 

6 - Key 哈希存储（SHA-256），不存明文 

7 - 前缀匹配快速查找 

8 - 权限范围（scope）绑定 

9 - 过期时间 & 用量配额 

10 - 审计日志联动 

11""" 

12 

13from __future__ import annotations 

14 

15import hashlib 

16import hmac 

17import secrets 

18import time 

19from dataclasses import dataclass, field 

20from enum import Enum 

21from typing import Optional 

22 

23 

24class KeyScope(str, Enum): 

25 """API Key 权限范围。""" 

26 READ = "read" 

27 WRITE = "write" 

28 ADMIN = "admin" 

29 AGENT_RUN = "agent:run" 

30 AGENT_MANAGE = "agent:manage" 

31 TOOLS_ALL = "tools:*" 

32 

33 

34@dataclass 

35class APIKey: 

36 """API Key 实体。只存储哈希，不存明文。""" 

37 key_id: str # 唯一标识，如 "ak_abc123" 

38 key_hash: str # SHA-256 哈希 

39 key_prefix: str # 前 8 位明文，用于快速匹配 

40 name: str # 人类可读名称，如 "生产环境 Bot" 

41 created_by: str # 创建者 

42 scopes: list[KeyScope] # 权限范围 

43 created_at: float = field(default_factory=time.time) 

44 expires_at: Optional[float] = None # 过期时间戳（None = 永不过期） 

45 last_used_at: Optional[float] = None 

46 usage_count: int = 0 

47 revoked: bool = False 

48 revoked_at: Optional[float] = None 

49 metadata: dict = field(default_factory=dict) 

50 

51 

52@dataclass 

53class KeyCreateRequest: 

54 """创建 API Key 的请求。""" 

55 name: str 

56 scopes: list[KeyScope] 

57 expires_in_days: Optional[int] = None # None = 永不过期 

58 metadata: dict = field(default_factory=dict) 

59 

60 

61@dataclass 

62class KeyCreateResult: 

63 """创建结果 — 包含仅此一次可见的明文 Key。""" 

64 key_id: str 

65 plaintext_key: str # ⚠️ 仅返回一次 

66 key_prefix: str 

67 scopes: list[KeyScope] 

68 expires_at: Optional[float] 

69 

70 

71class APIKeyManager: 

72 """API Key 全生命周期管理器。 

73 

74 特性： 

75 - SHA-256 哈希存储，不存明文 

76 - 前缀快速查找（前 8 位明文索引） 

77 - 撤销/轮转/用量追踪 

78 - 范围校验 

79 """ 

80 

81 def __init__(self, secret_salt: str = ""): 

82 self._keys: dict[str, APIKey] = {} # key_id → APIKey 

83 self._prefix_index: dict[str, str] = {} # key_prefix → key_id 

84 self._secret_salt = secret_salt or secrets.token_hex(16) 

85 

86 # ── 创建 ── 

87 

88 def create_key(self, request: KeyCreateRequest, created_by: str = "admin") -> KeyCreateResult: 

89 """创建新的 API Key。返回仅一次可见的明文。""" 

90 key_id = f"ak_{secrets.token_hex(12)}" 

91 plaintext = f"agentos_{secrets.token_hex(24)}" 

92 key_prefix = plaintext[:12] 

93 key_hash = self._hash(plaintext) 

94 

95 expires_at = None 

96 if request.expires_in_days: 

97 expires_at = time.time() + request.expires_in_days * 86400 

98 

99 key = APIKey( 

100 key_id=key_id, 

101 key_hash=key_hash, 

102 key_prefix=key_prefix, 

103 name=request.name, 

104 created_by=created_by, 

105 scopes=request.scopes, 

106 expires_at=expires_at, 

107 metadata=request.metadata, 

108 ) 

109 

110 self._keys[key_id] = key 

111 self._prefix_index[key_prefix] = key_id 

112 

113 return KeyCreateResult( 

114 key_id=key_id, 

115 plaintext_key=plaintext, 

116 key_prefix=key_prefix, 

117 scopes=request.scopes, 

118 expires_at=expires_at, 

119 ) 

120 

121 # ── 验证 ── 

122 

123 def validate_key(self, plaintext: str) -> Optional[APIKey]: 

124 """验证 API Key 并返回对应的 Key 对象。无效/已撤销/过期返回 None。""" 

125 key_hash = self._hash(plaintext) 

126 

127 # 前缀快速定位 

128 key_prefix = plaintext[:12] 

129 key_id = self._prefix_index.get(key_prefix) 

130 if not key_id: 

131 return None 

132 

133 key = self._keys.get(key_id) 

134 if not key: 

135 return None 

136 

137 # 恒定时间比对防时序攻击 

138 if not hmac.compare_digest(key.key_hash, key_hash): 

139 return None 

140 

141 if key.revoked: 

142 return None 

143 

144 if key.expires_at and time.time() > key.expires_at: 

145 return None 

146 

147 # 更新使用记录 

148 key.last_used_at = time.time() 

149 key.usage_count += 1 

150 

151 return key 

152 

153 def check_scope(self, key: APIKey, required_scope: KeyScope) -> bool: 

154 """检查 Key 是否拥有指定权限范围。""" 

155 if KeyScope.ADMIN in key.scopes: 

156 return True 

157 return required_scope in key.scopes 

158 

159 # ── 管理 ── 

160 

161 def revoke_key(self, key_id: str) -> bool: 

162 """撤销 API Key。""" 

163 key = self._keys.get(key_id) 

164 if not key or key.revoked: 

165 return False 

166 key.revoked = True 

167 key.revoked_at = time.time() 

168 return True 

169 

170 def rotate_key(self, key_id: str, created_by: str = "admin") -> Optional[KeyCreateResult]: 

171 """轮转 API Key：撤销旧 Key，创建新 Key。""" 

172 old = self._keys.get(key_id) 

173 if not old or old.revoked: 

174 return None 

175 

176 self.revoke_key(key_id) 

177 

178 expires_in = None 

179 if old.expires_at: 

180 remaining = old.expires_at - time.time() 

181 expires_in = int(max(1, remaining / 86400)) 

182 

183 return self.create_key( 

184 KeyCreateRequest( 

185 name=f"{old.name} (rotated)", 

186 scopes=old.scopes, 

187 expires_in_days=expires_in, 

188 metadata={"rotated_from": key_id, **old.metadata}, 

189 ), 

190 created_by=created_by, 

191 ) 

192 

193 def list_keys(self) -> list[APIKey]: 

194 """列出所有 Key（不含明文）。""" 

195 return sorted(self._keys.values(), key=lambda k: k.created_at, reverse=True) 

196 

197 def get_key(self, key_id: str) -> Optional[APIKey]: 

198 """获取单个 Key 信息。""" 

199 return self._keys.get(key_id) 

200 

201 def stats(self) -> dict: 

202 """Key 统计信息。""" 

203 total = len(self._keys) 

204 active = sum(1 for k in self._keys.values() if not k.revoked) 

205 revoked = total - active 

206 total_usage = sum(k.usage_count for k in self._keys.values()) 

207 return { 

208 "total": total, 

209 "active": active, 

210 "revoked": revoked, 

211 "total_usage_count": total_usage, 

212 } 

213 

214 # ── 内部 ── 

215 

216 def _hash(self, plaintext: str) -> str: 

217 """SHA-256 哈希（加盐）。""" 

218 return hashlib.sha256(f"{self._secret_salt}:{plaintext}".encode()).hexdigest()