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

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 StrEnum 

21 

22 

23class KeyScope(StrEnum): 

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

25 

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 

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

39 key_hash: str # SHA-256 哈希 

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

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

42 created_by: str # 创建者 

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

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

45 expires_at: float | None = None # 过期时间戳（None = 永不过期） 

46 last_used_at: float | None = None 

47 usage_count: int = 0 

48 revoked: bool = False 

49 revoked_at: float | None = None 

50 metadata: dict = field(default_factory=dict) 

51 

52 

53@dataclass 

54class KeyCreateRequest: 

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

56 

57 name: str 

58 scopes: list[KeyScope] 

59 expires_in_days: int | None = None # None = 永不过期 

60 metadata: dict = field(default_factory=dict) 

61 

62 

63@dataclass 

64class KeyCreateResult: 

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

66 

67 key_id: str 

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

69 key_prefix: str 

70 scopes: list[KeyScope] 

71 expires_at: float | None 

72 

73 

74class APIKeyManager: 

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

76 

77 特性： 

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

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

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

81 - 范围校验 

82 """ 

83 

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

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

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

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

88 

89 # ── 创建 ── 

90 

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

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

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

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

95 key_prefix = plaintext[:12] 

96 key_hash = self._hash(plaintext) 

97 

98 expires_at = None 

99 if request.expires_in_days: 

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

101 

102 key = APIKey( 

103 key_id=key_id, 

104 key_hash=key_hash, 

105 key_prefix=key_prefix, 

106 name=request.name, 

107 created_by=created_by, 

108 scopes=request.scopes, 

109 expires_at=expires_at, 

110 metadata=request.metadata, 

111 ) 

112 

113 self._keys[key_id] = key 

114 self._prefix_index[key_prefix] = key_id 

115 

116 return KeyCreateResult( 

117 key_id=key_id, 

118 plaintext_key=plaintext, 

119 key_prefix=key_prefix, 

120 scopes=request.scopes, 

121 expires_at=expires_at, 

122 ) 

123 

124 # ── 验证 ── 

125 

126 def validate_key(self, plaintext: str) -> APIKey | None: 

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

128 key_hash = self._hash(plaintext) 

129 

130 # 前缀快速定位 

131 key_prefix = plaintext[:12] 

132 key_id = self._prefix_index.get(key_prefix) 

133 if not key_id: 

134 return None 

135 

136 key = self._keys.get(key_id) 

137 if not key: 

138 return None 

139 

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

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

142 return None 

143 

144 if key.revoked: 

145 return None 

146 

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

148 return None 

149 

150 # 更新使用记录 

151 key.last_used_at = time.time() 

152 key.usage_count += 1 

153 

154 return key 

155 

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

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

158 if KeyScope.ADMIN in key.scopes: 

159 return True 

160 return required_scope in key.scopes 

161 

162 # ── 管理 ── 

163 

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

165 """撤销 API Key。""" 

166 key = self._keys.get(key_id) 

167 if not key or key.revoked: 

168 return False 

169 key.revoked = True 

170 key.revoked_at = time.time() 

171 return True 

172 

173 def rotate_key(self, key_id: str, created_by: str = "admin") -> KeyCreateResult | None: 

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

175 old = self._keys.get(key_id) 

176 if not old or old.revoked: 

177 return None 

178 

179 self.revoke_key(key_id) 

180 

181 expires_in = None 

182 if old.expires_at: 

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

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

185 

186 return self.create_key( 

187 KeyCreateRequest( 

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

189 scopes=old.scopes, 

190 expires_in_days=expires_in, 

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

192 ), 

193 created_by=created_by, 

194 ) 

195 

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

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

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

199 

200 def get_key(self, key_id: str) -> APIKey | None: 

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

202 return self._keys.get(key_id) 

203 

204 def stats(self) -> dict: 

205 """Key 统计信息。""" 

206 total = len(self._keys) 

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

208 revoked = total - active 

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

210 return { 

211 "total": total, 

212 "active": active, 

213 "revoked": revoked, 

214 "total_usage_count": total_usage, 

215 } 

216 

217 # ── 内部 ── 

218 

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

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

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