Forge
search
pythonb1bb0b24
1from __future__ import annotations
2
3import json
4from typing import Any
5
6from fastapi import HTTPException
7
8from open_agent_registry.auth import hash_secret, new_claim_code
9from open_agent_registry.channels import normalize_email, send_email_code, send_telegram_code
10from open_agent_registry.config import settings
11from open_agent_registry.db import Database, _utc_now
12from open_agent_registry.totp import new_totp_secret, otpauth_uri, verify_totp
13
14CHANNEL_EMAIL = "email"
15CHANNEL_TELEGRAM = "telegram"
16CHANNEL_TOTP = "totp"
17STEP_EMAIL_VERIFIED = "email_verified"
18
19
20def _get_pending_row(conn: Any, token: str) -> Any:
21 row = conn.execute("SELECT * FROM agents WHERE claim_token = ?", (token,)).fetchone()
22 if row is None:
23 raise HTTPException(status_code=404, detail="Invalid claim token")
24 if row["claim_status"] == "claimed":
25 raise HTTPException(status_code=409, detail="Already claimed")
26 return row
27
28
29def begin_claim(
30 db: Database,
31 token: str,
32 *,
33 email: str,
34 channel: str,
35 telegram_chat_id: str | None = None,
36) -> dict[str, str]:
37 owner_email = normalize_email(email)
38 channel = channel.strip().lower()
39 if channel not in {CHANNEL_EMAIL, CHANNEL_TELEGRAM, CHANNEL_TOTP}:
40 raise HTTPException(status_code=400, detail="channel must be email, telegram, or totp")
41
42 with db.connect() as conn:
43 row = _get_pending_row(conn, token)
44 agent_name = row["name"]
45 now = _utc_now()
46
47 if channel == CHANNEL_TOTP:
48 secret = new_totp_secret()
49 conn.execute(
50 """
51 UPDATE agents SET
52 owner_email = ?,
53 pending_claim_channel = ?,
54 pending_totp_secret = ?,
55 claim_code_hash = NULL,
56 claim_step = NULL,
57 owner_telegram_chat_id = ?,
58 updated_at = ?
59 WHERE claim_token = ?
60 """,
61 (
62 owner_email,
63 CHANNEL_TOTP,
64 secret,
65 telegram_chat_id,
66 now,
67 token,
68 ),
69 )
70 uri = otpauth_uri(secret, account_name=f"{agent_name}:{owner_email}")
71 payload = {
72 "channel": CHANNEL_TOTP,
73 "message": "Scan otpauth_uri in your authenticator app, then confirm with a 6-digit code.",
74 "otpauth_uri": uri,
75 "email": owner_email,
76 }
77 if settings.dev_expose_totp_secret:
78 payload["dev_totp_secret"] = secret
79 payload["note"] = "dev_totp_secret only when OAR_DEV_EXPOSE_TOTP_SECRET=true"
80 return payload
81
82 if channel == CHANNEL_TELEGRAM:
83 if not telegram_chat_id or not str(telegram_chat_id).strip():
84 raise HTTPException(status_code=400, detail="telegram_chat_id required for telegram channel")
85 if not settings.telegram_bot_token:
86 raise HTTPException(status_code=503, detail="Telegram bot not configured (OAR_TELEGRAM_BOT_TOKEN)")
87
88 code = new_claim_code()
89 conn.execute(
90 """
91 UPDATE agents SET
92 owner_email = ?,
93 pending_claim_channel = ?,
94 pending_totp_secret = NULL,
95 claim_code_hash = ?,
96 claim_step = NULL,
97 owner_telegram_chat_id = ?,
98 updated_at = ?
99 WHERE claim_token = ?
100 """,
101 (
102 owner_email,
103 channel,
104 hash_secret(code),
105 telegram_chat_id if channel == CHANNEL_TELEGRAM else None,
106 now,
107 token,
108 ),
109 )
110
111 payload: dict[str, str] = {
112 "channel": channel,
113 "message": "Verification code issued.",
114 "email": owner_email,
115 }
116
117 if channel == CHANNEL_EMAIL:
118 sent = send_email_code(owner_email, agent_name, code)
119 if sent:
120 payload["delivery"] = "smtp"
121 elif settings.dev_expose_claim_codes:
122 payload["dev_code"] = code
123 payload["delivery"] = "dev_json"
124 payload["note"] = "Configure OAR_SMTP_* for email delivery; dev_code when OAR_DEV_EXPOSE_CLAIM_CODES=true"
125 else:
126 raise HTTPException(
127 status_code=503,
128 detail="SMTP not configured and dev codes disabled",
129 )
130 elif channel == CHANNEL_TELEGRAM:
131 send_telegram_code(str(telegram_chat_id), agent_name, code)
132 payload["delivery"] = "telegram"
133 if settings.dev_expose_claim_codes:
134 payload["dev_code"] = code
135
136 return payload
137
138
139def begin_claim_2fa(db: Database, token: str, *, email: str) -> dict[str, str]:
140 """Step 1 of 2FA claim: email code only."""
141 owner_email = normalize_email(email)
142 code = new_claim_code()
143
144 with db.connect() as conn:
145 row = _get_pending_row(conn, token)
146 agent_name = row["name"]
147 conn.execute(
148 """
149 UPDATE agents SET
150 owner_email = ?,
151 pending_claim_channel = ?,
152 pending_totp_secret = NULL,
153 claim_code_hash = ?,
154 claim_step = NULL,
155 updated_at = ?
156 WHERE claim_token = ?
157 """,
158 (owner_email, CHANNEL_EMAIL, hash_secret(code), _utc_now(), token),
159 )
160
161 payload: dict[str, str] = {
162 "mode": "2fa",
163 "step": "1",
164 "next": "confirm-email then setup-totp",
165 "email": owner_email,
166 }
167 sent = send_email_code(owner_email, agent_name, code)
168 if sent:
169 payload["delivery"] = "smtp"
170 elif settings.dev_expose_claim_codes:
171 payload["dev_code"] = code
172 payload["delivery"] = "dev_json"
173 else:
174 raise HTTPException(status_code=503, detail="SMTP not configured and dev codes disabled")
175 return payload
176
177
178def setup_totp_2fa(db: Database, token: str) -> dict[str, str]:
179 """Step 2 of 2FA claim: enroll authenticator after email verified."""
180 with db.connect() as conn:
181 row = _get_pending_row(conn, token)
182 if row["claim_step"] != STEP_EMAIL_VERIFIED:
183 raise HTTPException(status_code=400, detail="Complete email verification first (POST .../confirm-email)")
184 secret = new_totp_secret()
185 conn.execute(
186 """
187 UPDATE agents SET pending_totp_secret = ?, pending_claim_channel = ?, updated_at = ?
188 WHERE claim_token = ?
189 """,
190 (secret, CHANNEL_TOTP, _utc_now(), token),
191 )
192 owner_email = row["owner_email"] or "owner"
193 uri = otpauth_uri(secret, account_name=f"{row['name']}:{owner_email}")
194
195 payload = {
196 "mode": "2fa",
197 "step": "2",
198 "message": "Scan otpauth_uri, then POST .../confirm-totp",
199 "otpauth_uri": uri,
200 }
201 if settings.dev_expose_totp_secret:
202 payload["dev_totp_secret"] = secret
203 return payload
204
205
206def confirm_claim(db: Database, token: str, *, email: str, code: str) -> dict[str, str]:
207 owner_email = normalize_email(email)
208 with db.connect() as conn:
209 row = _get_pending_row(conn, token)
210 if (row["owner_email"] or "").lower() != owner_email:
211 raise HTTPException(status_code=400, detail="Email does not match pending claim")
212
213 channel = row["pending_claim_channel"] or CHANNEL_EMAIL
214 if channel == CHANNEL_TOTP:
215 secret = row["pending_totp_secret"]
216 if not secret or not verify_totp(secret, code):
217 raise HTTPException(status_code=400, detail="Invalid authenticator code")
218 conn.execute(
219 """
220 UPDATE agents SET
221 claim_status = 'claimed',
222 owner_totp_secret = ?,
223 pending_totp_secret = NULL,
224 claim_code_hash = NULL,
225 claim_method = ?,
226 updated_at = ?
227 WHERE claim_token = ?
228 """,
229 (secret, CHANNEL_TOTP, _utc_now(), token),
230 )
231 return {"status": "claimed", "message": "Agent claimed with authenticator.", "owner_totp_enabled": "true"}
232
233 code_hash = hash_secret(code.strip())
234 if row["claim_code_hash"] != code_hash:
235 raise HTTPException(status_code=400, detail="Invalid verification code")
236
237 method = channel if channel in {CHANNEL_EMAIL, CHANNEL_TELEGRAM} else CHANNEL_EMAIL
238 conn.execute(
239 """
240 UPDATE agents SET
241 claim_status = 'claimed',
242 claim_code_hash = NULL,
243 claim_method = ?,
244 updated_at = ?
245 WHERE claim_token = ?
246 """,
247 (method, _utc_now(), token),
248 )
249 return {"status": "claimed", "message": f"Agent claimed via {method}."}
250
251
252def confirm_email_step(db: Database, token: str, *, email: str, code: str) -> dict[str, str]:
253 """2FA step 1 — verify email, then setup-totp."""
254 if not settings.claim_require_2fa:
255 return confirm_claim(db, token, email=email, code=code)
256
257 owner_email = normalize_email(email)
258 code_hash = hash_secret(code.strip())
259
260 with db.connect() as conn:
261 row = _get_pending_row(conn, token)
262 if (row["owner_email"] or "").lower() != owner_email:
263 raise HTTPException(status_code=400, detail="Email does not match pending claim")
264 if row["claim_code_hash"] != code_hash:
265 raise HTTPException(status_code=400, detail="Invalid verification code")
266
267 conn.execute(
268 """
269 UPDATE agents SET claim_step = ?, claim_code_hash = NULL, updated_at = ?
270 WHERE claim_token = ?
271 """,
272 (STEP_EMAIL_VERIFIED, _utc_now(), token),
273 )
274 return {
275 "status": "email_verified",
276 "next": f"POST /claim/{token}/setup-totp",
277 }
278
279
280def confirm_totp(db: Database, token: str, *, email: str, code: str) -> dict[str, str]:
281 owner_email = normalize_email(email)
282
283 with db.connect() as conn:
284 row = _get_pending_row(conn, token)
285 if (row["owner_email"] or "").lower() != owner_email:
286 raise HTTPException(status_code=400, detail="Email does not match pending claim")
287 secret = row["pending_totp_secret"]
288 if not secret:
289 raise HTTPException(status_code=400, detail="TOTP not started; POST .../begin with channel=totp")
290
291 if not verify_totp(secret, code):
292 raise HTTPException(status_code=400, detail="Invalid authenticator code")
293
294 two_fa = settings.claim_require_2fa or row["claim_step"] == STEP_EMAIL_VERIFIED
295 conn.execute(
296 """
297 UPDATE agents SET
298 claim_status = 'claimed',
299 claim_code_hash = NULL,
300 pending_totp_secret = NULL,
301 owner_totp_secret = ?,
302 claim_step = NULL,
303 pending_claim_channel = ?,
304 claim_method = ?,
305 updated_at = ?
306 WHERE claim_token = ?
307 """,
308 (
309 secret,
310 CHANNEL_TOTP,
311 "2fa" if two_fa else CHANNEL_TOTP,
312 _utc_now(),
313 token,
314 ),
315 )
316 method = "2fa" if two_fa else CHANNEL_TOTP
317 return {
318 "status": "claimed",
319 "message": "Agent claimed with authenticator.",
320 "owner_totp_enabled": "true",
321 "claim_method": method,
322 }
323
View only · write via MCP/CIDE