245 lines
8.3 KiB
Python
245 lines
8.3 KiB
Python
"""
|
|
Google OAuth handler Lambda.
|
|
|
|
Routes:
|
|
GET /oauth/start?actor_id=telegram:123&label=work → redirect to Google OAuth consent
|
|
GET /oauth/callback?code=...&state=... → exchange code, store tokens, update DynamoDB
|
|
"""
|
|
import base64
|
|
import json
|
|
import os
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
|
|
import boto3
|
|
|
|
_sm = None
|
|
_ddb = None
|
|
|
|
SCOPES = ' '.join([
|
|
'https://www.googleapis.com/auth/gmail.modify',
|
|
'https://www.googleapis.com/auth/calendar',
|
|
'https://www.googleapis.com/auth/drive',
|
|
'https://www.googleapis.com/auth/spreadsheets',
|
|
'https://www.googleapis.com/auth/documents',
|
|
'openid',
|
|
'email',
|
|
'profile',
|
|
])
|
|
|
|
|
|
def get_sm():
|
|
global _sm
|
|
if _sm is None:
|
|
_sm = boto3.client('secretsmanager', region_name=os.environ.get('AWS_REGION', 'us-east-1'))
|
|
return _sm
|
|
|
|
|
|
def get_ddb():
|
|
global _ddb
|
|
if _ddb is None:
|
|
_ddb = boto3.resource('dynamodb')
|
|
return _ddb
|
|
|
|
|
|
def get_oauth_client() -> tuple[str, str]:
|
|
"""Return (client_id, client_secret) from Secrets Manager."""
|
|
arn = os.environ['GOOGLE_OAUTH_CLIENT_SECRET_ARN']
|
|
secret = json.loads(get_sm().get_secret_value(SecretId=arn)['SecretString'])
|
|
return secret['client_id'], secret['client_secret']
|
|
|
|
|
|
def actor_id_to_secret_name(actor_id: str, label: str = 'primary') -> str:
|
|
safe = actor_id.replace(':', '-').replace('/', '-')
|
|
return f'agent-claw/google-credentials/{safe}/{label}'
|
|
|
|
|
|
def _redirect(url: str) -> dict:
|
|
return {'statusCode': 302, 'headers': {'Location': url}, 'body': ''}
|
|
|
|
|
|
def _html(body: str, status: int = 200) -> dict:
|
|
return {'statusCode': status, 'headers': {'Content-Type': 'text/html'}, 'body': body}
|
|
|
|
|
|
def handler(event, context):
|
|
path = event.get('rawPath') or event.get('path', '')
|
|
params = event.get('queryStringParameters') or {}
|
|
|
|
if path.endswith('/oauth/start'):
|
|
return handle_start(params)
|
|
elif path.endswith('/oauth/callback'):
|
|
return handle_callback(params)
|
|
else:
|
|
return {'statusCode': 404, 'body': 'Not found'}
|
|
|
|
|
|
def handle_start(params: dict) -> dict:
|
|
actor_id = params.get('actor_id', '')
|
|
if not actor_id:
|
|
return _html('<h1>Missing actor_id</h1>', 400)
|
|
|
|
label = params.get('label', 'primary')
|
|
|
|
client_id, _ = get_oauth_client()
|
|
redirect_uri = os.environ['OAUTH_REDIRECT_URI']
|
|
|
|
# Encode actor_id + label in state (JSON → base64)
|
|
state_data = json.dumps({'a': actor_id, 'l': label})
|
|
state = base64.urlsafe_b64encode(state_data.encode()).decode().rstrip('=')
|
|
|
|
auth_url = (
|
|
'https://accounts.google.com/o/oauth2/v2/auth?'
|
|
+ urllib.parse.urlencode({
|
|
'client_id': client_id,
|
|
'redirect_uri': redirect_uri,
|
|
'response_type': 'code',
|
|
'scope': SCOPES,
|
|
'access_type': 'offline',
|
|
'prompt': 'consent',
|
|
'state': state,
|
|
})
|
|
)
|
|
return _redirect(auth_url)
|
|
|
|
|
|
def handle_callback(params: dict) -> dict:
|
|
code = params.get('code', '')
|
|
state = params.get('state', '')
|
|
error = params.get('error', '')
|
|
|
|
if error:
|
|
return _html(f'<h1>OAuth error: {error}</h1>', 400)
|
|
if not code or not state:
|
|
return _html('<h1>Missing code or state</h1>', 400)
|
|
|
|
# Decode actor_id + label from state
|
|
try:
|
|
padding = 4 - len(state) % 4
|
|
state_data = json.loads(base64.urlsafe_b64decode(state + '=' * padding).decode())
|
|
actor_id = state_data['a']
|
|
label = state_data.get('l', 'primary')
|
|
except Exception:
|
|
# Backward compat: old state was just base64(actor_id)
|
|
try:
|
|
padding = 4 - len(state) % 4
|
|
actor_id = base64.urlsafe_b64decode(state + '=' * padding).decode()
|
|
label = 'primary'
|
|
except Exception:
|
|
return _html('<h1>Invalid state</h1>', 400)
|
|
|
|
client_id, client_secret = get_oauth_client()
|
|
redirect_uri = os.environ['OAUTH_REDIRECT_URI']
|
|
|
|
# Exchange code for tokens
|
|
token_data = urllib.parse.urlencode({
|
|
'code': code,
|
|
'client_id': client_id,
|
|
'client_secret': client_secret,
|
|
'redirect_uri': redirect_uri,
|
|
'grant_type': 'authorization_code',
|
|
}).encode()
|
|
|
|
req = urllib.request.Request(
|
|
'https://oauth2.googleapis.com/token',
|
|
data=token_data,
|
|
headers={'Content-Type': 'application/x-www-form-urlencoded'},
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=15) as resp:
|
|
tokens = json.loads(resp.read())
|
|
except Exception as e:
|
|
print(f'[oauth] Token exchange failed: {e}')
|
|
return _html(f'<h1>Token exchange failed: {e}</h1>', 500)
|
|
|
|
# Fetch user email from Google
|
|
user_email = ''
|
|
try:
|
|
id_token_payload = tokens.get('id_token', '').split('.')[1]
|
|
padding = 4 - len(id_token_payload) % 4
|
|
claims = json.loads(base64.urlsafe_b64decode(id_token_payload + '=' * padding))
|
|
user_email = claims.get('email', '')
|
|
except Exception:
|
|
pass
|
|
|
|
if not user_email:
|
|
try:
|
|
access_token = tokens.get('access_token', '')
|
|
req2 = urllib.request.Request(
|
|
'https://www.googleapis.com/oauth2/v3/userinfo',
|
|
headers={'Authorization': f'Bearer {access_token}'},
|
|
)
|
|
with urllib.request.urlopen(req2, timeout=10) as resp2:
|
|
user_email = json.loads(resp2.read()).get('email', '')
|
|
except Exception as e:
|
|
print(f'[oauth] userinfo fetch failed: {e}')
|
|
|
|
# Build credentials dict (google-auth format)
|
|
creds = {
|
|
'token': tokens.get('access_token'),
|
|
'refresh_token': tokens.get('refresh_token'),
|
|
'token_uri': 'https://oauth2.googleapis.com/token',
|
|
'client_id': client_id,
|
|
'client_secret': client_secret,
|
|
'scopes': [s for s in SCOPES.split() if s.startswith('https://')],
|
|
'email': user_email,
|
|
'user_email': user_email,
|
|
}
|
|
if tokens.get('expires_in'):
|
|
creds['expiry'] = time.strftime(
|
|
'%Y-%m-%dT%H:%M:%SZ',
|
|
time.gmtime(time.time() + int(tokens['expires_in']))
|
|
)
|
|
|
|
# Store in Secrets Manager at labeled path
|
|
secret_name = actor_id_to_secret_name(actor_id, label)
|
|
sm = get_sm()
|
|
try:
|
|
sm.create_secret(Name=secret_name, SecretString=json.dumps(creds))
|
|
except sm.exceptions.ResourceExistsException:
|
|
sm.put_secret_value(SecretId=secret_name, SecretString=json.dumps(creds))
|
|
print(f'[oauth] Stored credentials for actor={actor_id} label={label} email={user_email}')
|
|
|
|
# Update DynamoDB: merge into google_accounts map
|
|
table_name = os.environ.get('USERS_TABLE_NAME', '')
|
|
if table_name and actor_id:
|
|
try:
|
|
table = get_ddb().Table(table_name)
|
|
table.update_item(
|
|
Key={'actor_id': actor_id},
|
|
UpdateExpression='SET google_accounts = if_not_exists(google_accounts, :empty)',
|
|
ExpressionAttributeValues={':empty': {}},
|
|
)
|
|
table.update_item(
|
|
Key={'actor_id': actor_id},
|
|
UpdateExpression='SET google_accounts.#label = :email',
|
|
ExpressionAttributeNames={'#label': label},
|
|
ExpressionAttributeValues={':email': user_email},
|
|
)
|
|
except Exception as e:
|
|
print(f'[oauth] DynamoDB update failed: {e}')
|
|
|
|
# Best-effort Telegram confirmation
|
|
try:
|
|
bot_token_arn = os.environ.get('TELEGRAM_BOT_TOKEN_SECRET_ARN', '')
|
|
if bot_token_arn and actor_id.startswith('telegram:'):
|
|
chat_id = actor_id.split(':', 1)[1]
|
|
bot_token = get_sm().get_secret_value(SecretId=bot_token_arn)['SecretString']
|
|
tg_text = f'✅ Connected {user_email} as "{label}"'
|
|
tg_payload = json.dumps({'chat_id': chat_id, 'text': tg_text}).encode()
|
|
tg_req = urllib.request.Request(
|
|
f'https://api.telegram.org/bot{bot_token}/sendMessage',
|
|
data=tg_payload,
|
|
headers={'Content-Type': 'application/json'},
|
|
)
|
|
urllib.request.urlopen(tg_req, timeout=5)
|
|
except Exception as e:
|
|
print(f'[oauth] Telegram notification failed: {e}')
|
|
|
|
return _html(
|
|
f'<h1>✅ Google account connected!</h1>'
|
|
f'<p>Connected <b>{user_email}</b> as "<b>{label}</b>".</p>'
|
|
f'<p>You can close this window and return to Telegram.</p>'
|
|
)
|