agent-claw/cdk/cdk.out/asset.5be87975e51a6859dfad098b3d998a0bcd09a4f9a437bbf38923338fb559eb9e/handler.py

"""
Google OAuth handler Lambda.

Routes:
  GET /oauth/start?actor_id=telegram:123  → redirect to Google OAuth consent
  GET /oauth/callback?code=...&state=...  → exchange code, store tokens, update DynamoDB
"""
import base64
import hashlib
import hmac
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) -> str:
    safe = actor_id.replace(':', '-').replace('/', '-')
    return f'agent-claw/google-credentials/{safe}'


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)

    client_id, _ = get_oauth_client()
    redirect_uri = os.environ['OAUTH_REDIRECT_URI']

    # Encode actor_id in state (base64 to keep URL-safe)
    state = base64.urlsafe_b64encode(actor_id.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 from state
    try:
        padding = 4 - len(state) % 4
        actor_id = base64.urlsafe_b64decode(state + '=' * padding).decode()
    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:
        # Fallback: call userinfo endpoint
        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': SCOPES.split(),
        '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
    secret_name = actor_id_to_secret_name(actor_id)
    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} email={user_email}')

    # Update DynamoDB users table with google_email
    table_name = os.environ.get('USERS_TABLE_NAME', '')
    if table_name and actor_id:
        try:
            get_ddb().Table(table_name).update_item(
                Key={'actor_id': actor_id},
                UpdateExpression='SET google_email = :e',
                ExpressionAttributeValues={':e': user_email},
            )
        except Exception as e:
            print(f'[oauth] DynamoDB update failed: {e}')

    return _html(
        f'<h1>✅ Google account connected!</h1>'
        f'<p>Connected <b>{user_email}</b> to your agent account.</p>'
        f'<p>You can close this window and return to Telegram.</p>'
    )