|
| 1 | +""" |
| 2 | +MCP token authentication middleware for InfraBox. |
| 3 | +
|
| 4 | +Token format: ib_mcp_<48 hex chars> |
| 5 | +Lookup key: first 16 chars of the 48-char hex suffix |
| 6 | +Hash: SHA-256 of the full raw token string (UTF-8) |
| 7 | +
|
| 8 | +MCP tokens are ONLY accepted on /api/v1/mcp/* paths. |
| 9 | +All other paths fall back to the normal OPA-based auth. |
| 10 | +""" |
| 11 | +import hashlib |
| 12 | +import logging |
| 13 | +from datetime import datetime, timezone |
| 14 | +from functools import wraps |
| 15 | + |
| 16 | +from flask import g, request, abort |
| 17 | + |
| 18 | +logger = logging.getLogger('mcp_auth') |
| 19 | + |
| 20 | +_MCP_TOKEN_PREFIX = 'ib_mcp_' |
| 21 | +_MCP_PATH_PREFIX = '/api/v1/mcp/' |
| 22 | + |
| 23 | + |
| 24 | +def _utcnow(): |
| 25 | + return datetime.now(timezone.utc) |
| 26 | + |
| 27 | + |
| 28 | +def _utcnow_naive(): |
| 29 | + """Naive UTC datetime for comparing against psycopg2 TIMESTAMP values.""" |
| 30 | + return datetime.now(timezone.utc).replace(tzinfo=None) |
| 31 | + |
| 32 | + |
| 33 | +def _hash_token(raw_token: str) -> str: |
| 34 | + return hashlib.sha256(raw_token.encode('utf-8')).hexdigest() |
| 35 | + |
| 36 | + |
| 37 | +def _reject(status, message): |
| 38 | + abort(status, message) |
| 39 | + |
| 40 | + |
| 41 | +def mcp_auth_required(f): |
| 42 | + """Decorator that validates ib_mcp_* Bearer tokens. |
| 43 | +
|
| 44 | + Sets on flask.g: |
| 45 | + g.mcp_token_id – token_id (16-char prefix) |
| 46 | + g.mcp_token_user_id – user uuid who owns the token |
| 47 | + g.mcp_enabled_projects – dict {project_id: expires_at_iso_or_None} |
| 48 | + g.mcp_allow_trigger – bool |
| 49 | + """ |
| 50 | + @wraps(f) |
| 51 | + def decorated(*args, **kwargs): |
| 52 | + auth = request.headers.get('Authorization', '') |
| 53 | + |
| 54 | + if not auth.startswith('Bearer ' + _MCP_TOKEN_PREFIX): |
| 55 | + # not an MCP token — fall through to OPA auth (already done in before_request) |
| 56 | + return f(*args, **kwargs) |
| 57 | + |
| 58 | + raw_token = auth[len('Bearer '):] |
| 59 | + |
| 60 | + # MCP tokens are only valid on /api/v1/mcp/* paths |
| 61 | + if not request.path.startswith(_MCP_PATH_PREFIX): |
| 62 | + _reject(403, 'MCP token can only be used on /api/v1/mcp/* endpoints') |
| 63 | + |
| 64 | + token_suffix = raw_token[len(_MCP_TOKEN_PREFIX):] |
| 65 | + if len(token_suffix) != 48: |
| 66 | + _reject(401, 'invalid MCP token format') |
| 67 | + |
| 68 | + token_id = token_suffix[:16] |
| 69 | + token_hash = _hash_token(raw_token) |
| 70 | + |
| 71 | + row = g.db.execute_one_dict(''' |
| 72 | + SELECT token_id, user_id, enabled_projects, allow_trigger, expires_at, revoked_at |
| 73 | + FROM mcp_token |
| 74 | + WHERE token_id = %s AND token_hash = %s |
| 75 | + ''', [token_id, token_hash]) |
| 76 | + |
| 77 | + if not row: |
| 78 | + _reject(401, 'invalid or unknown MCP token') |
| 79 | + |
| 80 | + if row['revoked_at'] is not None: |
| 81 | + _reject(401, 'MCP token has been revoked') |
| 82 | + |
| 83 | + if row['expires_at'] < _utcnow_naive(): |
| 84 | + _reject(401, 'MCP token has expired') |
| 85 | + |
| 86 | + # Update last_used_at (best-effort, non-fatal) |
| 87 | + try: |
| 88 | + g.db.execute( |
| 89 | + 'UPDATE mcp_token SET last_used_at = NOW() WHERE token_id = %s', |
| 90 | + [token_id] |
| 91 | + ) |
| 92 | + g.db.commit() |
| 93 | + except Exception as exc: |
| 94 | + logger.warning('failed to update last_used_at: %s', exc) |
| 95 | + |
| 96 | + g.mcp_token_id = token_id |
| 97 | + g.mcp_token_user_id = str(row['user_id']) |
| 98 | + g.mcp_enabled_projects = row['enabled_projects'] or {} |
| 99 | + g.mcp_allow_trigger = bool(row['allow_trigger']) |
| 100 | + # Suppress OPA check for MCP-authed requests |
| 101 | + g.token = {'type': 'mcp', 'user': {'id': str(row['user_id']), 'role': 'user'}} |
| 102 | + |
| 103 | + return f(*args, **kwargs) |
| 104 | + return decorated |
| 105 | + |
| 106 | + |
| 107 | +def check_project_access_mcp(project_id: str) -> bool: |
| 108 | + """Return True if the current request may access project_id. |
| 109 | +
|
| 110 | + MCP token path: project must be in g.mcp_enabled_projects and not past |
| 111 | + its per-project expiry (if set). |
| 112 | + Session path: delegates to OPA (already checked in before_request). |
| 113 | + """ |
| 114 | + if not hasattr(g, 'mcp_enabled_projects'): |
| 115 | + # session / OPA path — access already verified |
| 116 | + return True |
| 117 | + |
| 118 | + enabled = g.mcp_enabled_projects |
| 119 | + if project_id not in enabled: |
| 120 | + return False |
| 121 | + |
| 122 | + per_project_expiry = enabled[project_id] |
| 123 | + if per_project_expiry: |
| 124 | + try: |
| 125 | + exp = datetime.fromisoformat(per_project_expiry) |
| 126 | + # fromisoformat() on a naive string (no UTC offset) returns a naive |
| 127 | + # datetime; compare against naive UTC to avoid TypeError. |
| 128 | + now = _utcnow_naive() if exp.tzinfo is None else _utcnow() |
| 129 | + if exp < now: |
| 130 | + return False |
| 131 | + except (ValueError, TypeError): |
| 132 | + # Malformed expiry — treat as expired rather than silently granting access. |
| 133 | + return False |
| 134 | + |
| 135 | + return True |
| 136 | + |
| 137 | + |
| 138 | +def check_trigger_access_mcp() -> bool: |
| 139 | + """Return True if the current request may trigger builds.""" |
| 140 | + if not hasattr(g, 'mcp_allow_trigger'): |
| 141 | + return True |
| 142 | + return bool(g.mcp_allow_trigger) |
| 143 | + |
| 144 | + |
| 145 | +def get_mcp_user_id() -> str: |
| 146 | + """Return user id string regardless of auth path.""" |
| 147 | + if hasattr(g, 'mcp_token_user_id'): |
| 148 | + return g.mcp_token_user_id |
| 149 | + token = getattr(g, 'token', None) |
| 150 | + if token and 'user' in token: |
| 151 | + return str(token['user'].get('id', '')) |
| 152 | + return request.remote_addr or 'unknown' |
0 commit comments