491 lines
17 KiB
Python
491 lines
17 KiB
Python
import asyncio
|
|
import logging
|
|
from functools import wraps
|
|
from uuid import UUID
|
|
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
|
from webauthn.helpers.exceptions import InvalidAuthenticationResponse
|
|
|
|
from paskia import remoteauth
|
|
from paskia.authsession import create_session, get_reset, get_session
|
|
from paskia.fastapi import authz
|
|
from paskia.fastapi.session import AUTH_COOKIE, infodict
|
|
from paskia.globals import db, passkey
|
|
from paskia.util import hostutil, passphrase
|
|
from paskia.util.tokens import create_token, session_key
|
|
|
|
|
|
# WebSocket error handling decorator
|
|
def websocket_error_handler(func):
|
|
@wraps(func)
|
|
async def wrapper(ws: WebSocket, *args, **kwargs):
|
|
try:
|
|
await ws.accept()
|
|
return await func(ws, *args, **kwargs)
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except authz.AuthException as e:
|
|
await ws.send_json(
|
|
{
|
|
"status": e.status_code,
|
|
**(await authz.auth_error_content(e)),
|
|
}
|
|
)
|
|
except (ValueError, InvalidAuthenticationResponse) as e:
|
|
await ws.send_json({"status": 401, "detail": str(e)})
|
|
except Exception:
|
|
logging.exception("Internal Server Error")
|
|
await ws.send_json({"status": 500, "detail": "Internal Server Error"})
|
|
|
|
return wrapper
|
|
|
|
|
|
# Create a FastAPI subapp for WebSocket endpoints
|
|
app = FastAPI()
|
|
|
|
|
|
def _validate_origin(ws: WebSocket) -> str:
|
|
"""Extract and validate origin from WebSocket request headers.
|
|
|
|
Raises:
|
|
ValueError: If origin header is missing or not in allowed list
|
|
"""
|
|
origin = ws.headers.get("origin")
|
|
if not origin:
|
|
raise ValueError("Origin header is required for WebSocket connections")
|
|
return passkey.instance.validate_origin(origin)
|
|
|
|
|
|
async def register_chat(
|
|
ws: WebSocket,
|
|
user_uuid: UUID,
|
|
user_name: str,
|
|
origin: str,
|
|
credential_ids: list[bytes] | None = None,
|
|
):
|
|
"""Generate registration options and send them to the client."""
|
|
options, challenge = passkey.instance.reg_generate_options(
|
|
user_id=user_uuid,
|
|
user_name=user_name,
|
|
credential_ids=credential_ids,
|
|
)
|
|
await ws.send_json({"optionsJSON": options})
|
|
response = await ws.receive_json()
|
|
return passkey.instance.reg_verify(response, challenge, user_uuid, origin=origin)
|
|
|
|
|
|
@app.websocket("/register")
|
|
@websocket_error_handler
|
|
async def websocket_register_add(
|
|
ws: WebSocket,
|
|
reset: str | None = None,
|
|
name: str | None = None,
|
|
auth=AUTH_COOKIE,
|
|
):
|
|
"""Register a new credential for an existing user.
|
|
|
|
Supports either:
|
|
- Normal session via auth cookie (requires recent authentication)
|
|
- Reset token supplied as ?reset=... (auth cookie ignored)
|
|
"""
|
|
origin = _validate_origin(ws)
|
|
host = origin.split("://", 1)[1]
|
|
if reset is not None:
|
|
if not passphrase.is_well_formed(reset):
|
|
raise ValueError(
|
|
f"The reset link for {passkey.instance.rp_name} is invalid or has expired"
|
|
)
|
|
s = await get_reset(reset)
|
|
user_uuid = s.user_uuid
|
|
else:
|
|
# Require recent authentication for adding a new passkey
|
|
ctx = await authz.verify(auth, perm=[], host=host, max_age="5m")
|
|
user_uuid = ctx.session.user_uuid
|
|
s = ctx.session
|
|
|
|
# Get user information and determine effective user_name for this registration
|
|
user = await db.instance.get_user_by_uuid(user_uuid)
|
|
user_name = user.display_name
|
|
if name is not None:
|
|
stripped = name.strip()
|
|
if stripped:
|
|
user_name = stripped
|
|
challenge_ids = await db.instance.get_credentials_by_user_uuid(user_uuid)
|
|
|
|
# WebAuthn registration
|
|
credential = await register_chat(ws, user_uuid, user_name, origin, challenge_ids)
|
|
|
|
# Create a new session and store everything in database
|
|
token = create_token()
|
|
metadata = infodict(ws, "authenticated")
|
|
await db.instance.create_credential_session( # type: ignore[attr-defined]
|
|
user_uuid=user_uuid,
|
|
credential=credential,
|
|
reset_key=(s.key if reset is not None else None),
|
|
session_key=session_key(token),
|
|
display_name=user_name,
|
|
host=host,
|
|
ip=metadata.get("ip"),
|
|
user_agent=metadata.get("user_agent"),
|
|
)
|
|
auth = token
|
|
|
|
assert isinstance(auth, str) and len(auth) == 16
|
|
await ws.send_json(
|
|
{
|
|
"user_uuid": str(user.uuid),
|
|
"credential_uuid": str(credential.uuid),
|
|
"session_token": auth,
|
|
"message": "New credential added successfully",
|
|
}
|
|
)
|
|
|
|
|
|
@app.websocket("/authenticate")
|
|
@websocket_error_handler
|
|
async def websocket_authenticate(ws: WebSocket, auth=AUTH_COOKIE):
|
|
origin = _validate_origin(ws)
|
|
host = origin.split("://", 1)[1]
|
|
|
|
# If there's an existing session, restrict to that user's credentials (reauth)
|
|
session_user_uuid = None
|
|
credential_ids = None
|
|
if auth:
|
|
try:
|
|
session = await get_session(auth, host=host)
|
|
session_user_uuid = session.user_uuid
|
|
credential_ids = await db.instance.get_credentials_by_user_uuid(
|
|
session_user_uuid
|
|
)
|
|
except ValueError:
|
|
pass # Invalid/expired session - allow normal authentication
|
|
|
|
options, challenge = passkey.instance.auth_generate_options(
|
|
credential_ids=credential_ids
|
|
)
|
|
await ws.send_json({"optionsJSON": options})
|
|
# Wait for the client to use his authenticator to authenticate
|
|
credential = passkey.instance.auth_parse(await ws.receive_json())
|
|
# Fetch from the database by credential ID
|
|
try:
|
|
stored_cred = await db.instance.get_credential_by_id(credential.raw_id)
|
|
except ValueError:
|
|
raise ValueError(
|
|
f"This passkey is no longer registered with {passkey.instance.rp_name}"
|
|
)
|
|
|
|
# If reauth mode, verify the credential belongs to the session's user
|
|
if session_user_uuid and stored_cred.user_uuid != session_user_uuid:
|
|
raise ValueError("This passkey belongs to a different account")
|
|
|
|
# Verify the credential matches the stored data
|
|
passkey.instance.auth_verify(credential, challenge, stored_cred, origin)
|
|
# Update both credential and user's last_seen timestamp
|
|
await db.instance.login(stored_cred.user_uuid, stored_cred)
|
|
|
|
# Create a session token for the authenticated user
|
|
assert stored_cred.uuid is not None
|
|
metadata = infodict(ws, "auth")
|
|
token = await create_session(
|
|
user_uuid=stored_cred.user_uuid,
|
|
credential_uuid=stored_cred.uuid,
|
|
host=host,
|
|
ip=metadata.get("ip") or "",
|
|
user_agent=metadata.get("user_agent") or "",
|
|
)
|
|
|
|
await ws.send_json(
|
|
{
|
|
"user_uuid": str(stored_cred.user_uuid),
|
|
"session_token": token,
|
|
}
|
|
)
|
|
|
|
|
|
@app.websocket("/remote-auth/request")
|
|
@websocket_error_handler
|
|
async def websocket_remote_auth_request(ws: WebSocket):
|
|
"""Request authentication from another device.
|
|
|
|
This endpoint is called by the device that wants to be authenticated.
|
|
It creates a remote auth request and waits for another device to authenticate.
|
|
|
|
Flow:
|
|
1. Client connects
|
|
2. Server creates a remote auth token and sends it with URL/expiry/pairing_code
|
|
3. Server waits for another device to authenticate via /remote-auth/complete
|
|
4. When auth completes, server sends session_token to this client
|
|
5. Client can then use the session token to set a cookie
|
|
"""
|
|
origin = _validate_origin(ws)
|
|
host = origin.split("://", 1)[1]
|
|
|
|
if remoteauth.instance is None:
|
|
raise ValueError("Remote authentication is not available")
|
|
|
|
metadata = infodict(ws, "remote-auth-request")
|
|
|
|
# Create the remote auth request
|
|
token, pairing_code, expiry = await remoteauth.instance.create_request(
|
|
host=host,
|
|
ip=metadata.get("ip") or "",
|
|
user_agent=metadata.get("user_agent") or "",
|
|
)
|
|
|
|
# Build the URL for the authenticating device (same endpoint as reset tokens)
|
|
url = hostutil.auth_site_base_url() + token
|
|
|
|
# Send the token, pairing code, and URL to the client
|
|
await ws.send_json(
|
|
{
|
|
"token": token,
|
|
"pairing_code": pairing_code,
|
|
"url": url,
|
|
"expires": expiry.isoformat().replace("+00:00", "Z"),
|
|
}
|
|
)
|
|
|
|
# Set up async notification
|
|
result_event = asyncio.Event()
|
|
result_data: dict = {}
|
|
|
|
def on_complete(
|
|
session_token: str | None,
|
|
user_uuid: UUID | None,
|
|
credential_uuid: UUID | None,
|
|
):
|
|
result_data["session_token"] = session_token
|
|
result_data["user_uuid"] = user_uuid
|
|
result_data["credential_uuid"] = credential_uuid
|
|
result_event.set()
|
|
|
|
await remoteauth.instance.set_notify_callback(token, on_complete)
|
|
|
|
try:
|
|
# Wait for either:
|
|
# 1. Authentication to complete (result_event set)
|
|
# 2. Client to disconnect
|
|
# 3. Client to send a cancel message
|
|
# 4. Timeout (handled by remoteauth cleanup)
|
|
|
|
while True:
|
|
# Use asyncio.wait to handle both event and websocket
|
|
receive_task = asyncio.create_task(ws.receive_json())
|
|
event_task = asyncio.create_task(result_event.wait())
|
|
|
|
done, pending = await asyncio.wait(
|
|
[receive_task, event_task],
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
|
|
# Cancel pending tasks
|
|
for task in pending:
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
if event_task in done:
|
|
# Authentication completed (or expired/cancelled)
|
|
if result_data.get("session_token"):
|
|
await ws.send_json(
|
|
{
|
|
"status": "authenticated",
|
|
"user_uuid": str(result_data["user_uuid"]),
|
|
"session_token": result_data["session_token"],
|
|
}
|
|
)
|
|
else:
|
|
await ws.send_json(
|
|
{
|
|
"status": "expired",
|
|
"detail": "Remote authentication request expired or was cancelled",
|
|
}
|
|
)
|
|
break
|
|
|
|
if receive_task in done:
|
|
# Client sent a message
|
|
msg = receive_task.result()
|
|
if msg.get("action") == "cancel":
|
|
await remoteauth.instance.cancel_request(token)
|
|
await ws.send_json({"status": "cancelled"})
|
|
break
|
|
# Ignore other messages
|
|
|
|
except WebSocketDisconnect:
|
|
# Client disconnected, cancel the request
|
|
await remoteauth.instance.cancel_request(token)
|
|
except Exception:
|
|
await remoteauth.instance.cancel_request(token)
|
|
raise
|
|
|
|
|
|
@app.websocket("/remote-auth/complete/{token}")
|
|
@websocket_error_handler
|
|
async def websocket_remote_auth_complete(ws: WebSocket, token: str):
|
|
"""Complete a remote authentication request.
|
|
|
|
This endpoint is called by the authenticating device (the one with the passkey).
|
|
It performs WebAuthn authentication and notifies the requesting device.
|
|
|
|
Flow:
|
|
1. Client opens the remote auth link and connects here
|
|
2. Server verifies the token is valid
|
|
3. Server sends WebAuthn options
|
|
4. Client authenticates with passkey
|
|
5. Server creates session for the REQUESTING device's host
|
|
6. Server notifies the requesting device via the callback
|
|
7. Server sends confirmation to this client
|
|
"""
|
|
origin = _validate_origin(ws)
|
|
|
|
if remoteauth.instance is None:
|
|
raise ValueError("Remote authentication is not available")
|
|
|
|
# Validate the remote auth token
|
|
request = await remoteauth.instance.get_request(token)
|
|
if request is None:
|
|
raise ValueError("This remote authentication link is invalid or has expired")
|
|
|
|
if request.completed:
|
|
raise ValueError("This remote authentication has already been completed")
|
|
|
|
# The session will be created for the requesting device's host, not this device's
|
|
target_host = request.host
|
|
|
|
# Generate authentication options (no credential restriction for remote auth)
|
|
options, challenge = passkey.instance.auth_generate_options(credential_ids=None)
|
|
await ws.send_json({"optionsJSON": options})
|
|
|
|
# Wait for client authentication response
|
|
credential = passkey.instance.auth_parse(await ws.receive_json())
|
|
|
|
# Fetch and verify credential
|
|
try:
|
|
stored_cred = await db.instance.get_credential_by_id(credential.raw_id)
|
|
except ValueError:
|
|
raise ValueError(
|
|
f"This passkey is no longer registered with {passkey.instance.rp_name}"
|
|
)
|
|
|
|
# Verify the credential
|
|
passkey.instance.auth_verify(credential, challenge, stored_cred, origin)
|
|
|
|
# Update credential last_used
|
|
await db.instance.login(stored_cred.user_uuid, stored_cred)
|
|
|
|
# Create a session for the REQUESTING device
|
|
assert stored_cred.uuid is not None
|
|
session_token = await create_session(
|
|
user_uuid=stored_cred.user_uuid,
|
|
credential_uuid=stored_cred.uuid,
|
|
host=target_host,
|
|
ip=request.ip,
|
|
user_agent=request.user_agent,
|
|
)
|
|
|
|
# Complete the remote auth request (notifies the waiting device)
|
|
completed = await remoteauth.instance.complete_request(
|
|
token=token,
|
|
session_token=session_token,
|
|
user_uuid=stored_cred.user_uuid,
|
|
credential_uuid=stored_cred.uuid,
|
|
)
|
|
|
|
if not completed:
|
|
raise ValueError("Failed to complete remote authentication")
|
|
|
|
# Send confirmation to the authenticating device
|
|
await ws.send_json(
|
|
{
|
|
"status": "success",
|
|
"message": "Authentication successful. The other device is now logged in.",
|
|
}
|
|
)
|
|
|
|
|
|
@app.websocket("/remote-auth/pair/{code}")
|
|
@websocket_error_handler
|
|
async def websocket_remote_auth_pair(ws: WebSocket, code: str):
|
|
"""Complete a remote authentication request using a pairing code.
|
|
|
|
This endpoint is called from the user's profile on the authenticating device.
|
|
The user enters the pairing code displayed on the requesting device.
|
|
|
|
Flow:
|
|
1. User on Device B (with passkey) enters pairing code from Device A
|
|
2. Server looks up the remote auth request by pairing code
|
|
3. Server sends WebAuthn options
|
|
4. User authenticates with passkey
|
|
5. Server creates session for Device A's host with Device A's metadata
|
|
6. Server notifies Device A via the callback
|
|
7. Server sends confirmation to Device B
|
|
"""
|
|
origin = _validate_origin(ws)
|
|
|
|
if remoteauth.instance is None:
|
|
raise ValueError("Remote authentication is not available")
|
|
|
|
# Look up the remote auth request by pairing code
|
|
request = await remoteauth.instance.get_request_by_pairing_code(code)
|
|
if request is None:
|
|
raise ValueError("Invalid or expired pairing code")
|
|
|
|
if request.completed:
|
|
raise ValueError("This remote authentication has already been completed")
|
|
|
|
# The session will be created for the requesting device's host
|
|
target_host = request.host
|
|
|
|
# Generate authentication options (no credential restriction for remote auth)
|
|
options, challenge = passkey.instance.auth_generate_options(credential_ids=None)
|
|
await ws.send_json({"optionsJSON": options})
|
|
|
|
# Wait for client authentication response
|
|
credential = passkey.instance.auth_parse(await ws.receive_json())
|
|
|
|
# Fetch and verify credential
|
|
try:
|
|
stored_cred = await db.instance.get_credential_by_id(credential.raw_id)
|
|
except ValueError:
|
|
raise ValueError(
|
|
f"This passkey is no longer registered with {passkey.instance.rp_name}"
|
|
)
|
|
|
|
# Verify the credential
|
|
passkey.instance.auth_verify(credential, challenge, stored_cred, origin)
|
|
|
|
# Update credential last_used
|
|
await db.instance.login(stored_cred.user_uuid, stored_cred)
|
|
|
|
# Create a session for the REQUESTING device (with their IP/user-agent)
|
|
assert stored_cred.uuid is not None
|
|
session_token = await create_session(
|
|
user_uuid=stored_cred.user_uuid,
|
|
credential_uuid=stored_cred.uuid,
|
|
host=target_host,
|
|
ip=request.ip,
|
|
user_agent=request.user_agent,
|
|
)
|
|
|
|
# Complete the remote auth request (notifies the waiting device)
|
|
completed = await remoteauth.instance.complete_request(
|
|
token=request.key,
|
|
session_token=session_token,
|
|
user_uuid=stored_cred.user_uuid,
|
|
credential_uuid=stored_cred.uuid,
|
|
)
|
|
|
|
if not completed:
|
|
raise ValueError("Failed to complete remote authentication")
|
|
|
|
# Send confirmation to the authenticating device
|
|
await ws.send_json(
|
|
{
|
|
"status": "success",
|
|
"message": "Authentication successful. The other device is now logged in.",
|
|
}
|
|
)
|