initial implementation
This commit is contained in:
parent
8f89000277
commit
7f324efd87
|
|
@ -0,0 +1,30 @@
|
||||||
|
__pycache__
|
||||||
|
*.pyc
|
||||||
|
*.pyo
|
||||||
|
*.pyd
|
||||||
|
.Python
|
||||||
|
*.so
|
||||||
|
*.egg
|
||||||
|
*.egg-info
|
||||||
|
dist
|
||||||
|
build
|
||||||
|
.env
|
||||||
|
.env.*
|
||||||
|
!.env.example
|
||||||
|
.venv
|
||||||
|
venv
|
||||||
|
ENV
|
||||||
|
env
|
||||||
|
.pytest_cache
|
||||||
|
.mypy_cache
|
||||||
|
.ruff_cache
|
||||||
|
.coverage
|
||||||
|
htmlcov
|
||||||
|
*.log
|
||||||
|
.git
|
||||||
|
.gitignore
|
||||||
|
README.md
|
||||||
|
*.md
|
||||||
|
tests
|
||||||
|
.vscode
|
||||||
|
.idea
|
||||||
|
|
@ -0,0 +1,46 @@
|
||||||
|
[style]
|
||||||
|
based_on_style = pep8
|
||||||
|
column_limit = 82
|
||||||
|
indent_width = 4
|
||||||
|
continuation_indent_width = 4
|
||||||
|
indent_closing_brackets = false
|
||||||
|
dedent_closing_brackets = true
|
||||||
|
indent_blank_lines = false
|
||||||
|
spaces_before_comment = 2
|
||||||
|
spaces_around_power_operator = false
|
||||||
|
spaces_around_default_or_named_assign = true
|
||||||
|
space_between_ending_comma_and_closing_bracket = false
|
||||||
|
space_inside_brackets = false
|
||||||
|
spaces_around_subscript_colon = true
|
||||||
|
blank_line_before_nested_class_or_def = false
|
||||||
|
blank_line_before_class_docstring = false
|
||||||
|
blank_lines_around_top_level_definition = 2
|
||||||
|
blank_lines_between_top_level_imports_and_variables = 2
|
||||||
|
blank_line_before_module_docstring = false
|
||||||
|
split_before_logical_operator = true
|
||||||
|
split_before_first_argument = true
|
||||||
|
split_before_named_assigns = true
|
||||||
|
split_complex_comprehension = true
|
||||||
|
split_before_expression_after_opening_paren = false
|
||||||
|
split_before_closing_bracket = true
|
||||||
|
split_all_comma_separated_values = true
|
||||||
|
split_all_top_level_comma_separated_values = false
|
||||||
|
coalesce_brackets = false
|
||||||
|
each_dict_entry_on_separate_line = true
|
||||||
|
allow_multiline_lambdas = false
|
||||||
|
allow_multiline_dictionary_keys = false
|
||||||
|
split_penalty_import_names = 0
|
||||||
|
join_multiple_lines = false
|
||||||
|
align_closing_bracket_with_visual_indent = true
|
||||||
|
arithmetic_precedence_indication = false
|
||||||
|
split_penalty_for_added_line_split = 275
|
||||||
|
use_tabs = false
|
||||||
|
split_before_dot = false
|
||||||
|
split_arguments_when_comma_terminated = true
|
||||||
|
i18n_function_call = ['_', 'N_', 'gettext', 'ngettext']
|
||||||
|
i18n_comment = ['# Translators:', '# i18n:']
|
||||||
|
split_penalty_comprehension = 80
|
||||||
|
split_penalty_after_opening_bracket = 280
|
||||||
|
split_penalty_before_if_expr = 0
|
||||||
|
split_penalty_bitwise_operator = 290
|
||||||
|
split_penalty_logical_operator = 0
|
||||||
|
|
@ -0,0 +1,147 @@
|
||||||
|
# A generic, single database configuration.
|
||||||
|
|
||||||
|
[alembic]
|
||||||
|
# path to migration scripts.
|
||||||
|
# this is typically a path given in POSIX (e.g. forward slashes)
|
||||||
|
# format, relative to the token %(here)s which refers to the location of this
|
||||||
|
# ini file
|
||||||
|
script_location = %(here)s/alembic
|
||||||
|
|
||||||
|
# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
|
||||||
|
# Uncomment the line below if you want the files to be prepended with date and time
|
||||||
|
# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
|
||||||
|
# for all available tokens
|
||||||
|
# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
|
||||||
|
|
||||||
|
# sys.path path, will be prepended to sys.path if present.
|
||||||
|
# defaults to the current working directory. for multiple paths, the path separator
|
||||||
|
# is defined by "path_separator" below.
|
||||||
|
prepend_sys_path = .
|
||||||
|
|
||||||
|
|
||||||
|
# timezone to use when rendering the date within the migration file
|
||||||
|
# as well as the filename.
|
||||||
|
# If specified, requires the tzdata library which can be installed by adding
|
||||||
|
# `alembic[tz]` to the pip requirements.
|
||||||
|
# string value is passed to ZoneInfo()
|
||||||
|
# leave blank for localtime
|
||||||
|
# timezone =
|
||||||
|
|
||||||
|
# max length of characters to apply to the "slug" field
|
||||||
|
# truncate_slug_length = 40
|
||||||
|
|
||||||
|
# set to 'true' to run the environment during
|
||||||
|
# the 'revision' command, regardless of autogenerate
|
||||||
|
# revision_environment = false
|
||||||
|
|
||||||
|
# set to 'true' to allow .pyc and .pyo files without
|
||||||
|
# a source .py file to be detected as revisions in the
|
||||||
|
# versions/ directory
|
||||||
|
# sourceless = false
|
||||||
|
|
||||||
|
# version location specification; This defaults
|
||||||
|
# to <script_location>/versions. When using multiple version
|
||||||
|
# directories, initial revisions must be specified with --version-path.
|
||||||
|
# The path separator used here should be the separator specified by "path_separator"
|
||||||
|
# below.
|
||||||
|
# version_locations = %(here)s/bar:%(here)s/bat:%(here)s/alembic/versions
|
||||||
|
|
||||||
|
# path_separator; This indicates what character is used to split lists of file
|
||||||
|
# paths, including version_locations and prepend_sys_path within configparser
|
||||||
|
# files such as alembic.ini.
|
||||||
|
# The default rendered in new alembic.ini files is "os", which uses os.pathsep
|
||||||
|
# to provide os-dependent path splitting.
|
||||||
|
#
|
||||||
|
# Note that in order to support legacy alembic.ini files, this default does NOT
|
||||||
|
# take place if path_separator is not present in alembic.ini. If this
|
||||||
|
# option is omitted entirely, fallback logic is as follows:
|
||||||
|
#
|
||||||
|
# 1. Parsing of the version_locations option falls back to using the legacy
|
||||||
|
# "version_path_separator" key, which if absent then falls back to the legacy
|
||||||
|
# behavior of splitting on spaces and/or commas.
|
||||||
|
# 2. Parsing of the prepend_sys_path option falls back to the legacy
|
||||||
|
# behavior of splitting on spaces, commas, or colons.
|
||||||
|
#
|
||||||
|
# Valid values for path_separator are:
|
||||||
|
#
|
||||||
|
# path_separator = :
|
||||||
|
# path_separator = ;
|
||||||
|
# path_separator = space
|
||||||
|
# path_separator = newline
|
||||||
|
#
|
||||||
|
# Use os.pathsep. Default configuration used for new projects.
|
||||||
|
path_separator = os
|
||||||
|
|
||||||
|
# set to 'true' to search source files recursively
|
||||||
|
# in each "version_locations" directory
|
||||||
|
# new in Alembic version 1.10
|
||||||
|
# recursive_version_locations = false
|
||||||
|
|
||||||
|
# the output encoding used when revision files
|
||||||
|
# are written from script.py.mako
|
||||||
|
# output_encoding = utf-8
|
||||||
|
|
||||||
|
# database URL. This is consumed by the user-maintained env.py script only.
|
||||||
|
# other means of configuring database URLs may be customized within the env.py
|
||||||
|
# file.
|
||||||
|
# sqlalchemy.url = driver://user:pass@localhost/dbname
|
||||||
|
|
||||||
|
|
||||||
|
[post_write_hooks]
|
||||||
|
# post_write_hooks defines scripts or Python functions that are run
|
||||||
|
# on newly generated revision scripts. See the documentation for further
|
||||||
|
# detail and examples
|
||||||
|
|
||||||
|
# format using "black" - use the console_scripts runner, against the "black" entrypoint
|
||||||
|
# hooks = black
|
||||||
|
# black.type = console_scripts
|
||||||
|
# black.entrypoint = black
|
||||||
|
# black.options = -l 79 REVISION_SCRIPT_FILENAME
|
||||||
|
|
||||||
|
# lint with attempts to fix using "ruff" - use the module runner, against the "ruff" module
|
||||||
|
# hooks = ruff
|
||||||
|
# ruff.type = module
|
||||||
|
# ruff.module = ruff
|
||||||
|
# ruff.options = check --fix REVISION_SCRIPT_FILENAME
|
||||||
|
|
||||||
|
# Alternatively, use the exec runner to execute a binary found on your PATH
|
||||||
|
# hooks = ruff
|
||||||
|
# ruff.type = exec
|
||||||
|
# ruff.executable = ruff
|
||||||
|
# ruff.options = check --fix REVISION_SCRIPT_FILENAME
|
||||||
|
|
||||||
|
# Logging configuration. This is also consumed by the user-maintained
|
||||||
|
# env.py script only.
|
||||||
|
[loggers]
|
||||||
|
keys = root,sqlalchemy,alembic
|
||||||
|
|
||||||
|
[handlers]
|
||||||
|
keys = console
|
||||||
|
|
||||||
|
[formatters]
|
||||||
|
keys = generic
|
||||||
|
|
||||||
|
[logger_root]
|
||||||
|
level = WARNING
|
||||||
|
handlers = console
|
||||||
|
qualname =
|
||||||
|
|
||||||
|
[logger_sqlalchemy]
|
||||||
|
level = WARNING
|
||||||
|
handlers =
|
||||||
|
qualname = sqlalchemy.engine
|
||||||
|
|
||||||
|
[logger_alembic]
|
||||||
|
level = INFO
|
||||||
|
handlers =
|
||||||
|
qualname = alembic
|
||||||
|
|
||||||
|
[handler_console]
|
||||||
|
class = StreamHandler
|
||||||
|
args = (sys.stderr,)
|
||||||
|
level = NOTSET
|
||||||
|
formatter = generic
|
||||||
|
|
||||||
|
[formatter_generic]
|
||||||
|
format = %(levelname)-5.5s [%(name)s] %(message)s
|
||||||
|
datefmt = %H:%M:%S
|
||||||
|
|
@ -0,0 +1 @@
|
||||||
|
Generic single database configuration
|
||||||
|
|
@ -0,0 +1,73 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Alembic environment configuration for SQLModel migrations
|
||||||
|
"""
|
||||||
|
|
||||||
|
from logging.config import fileConfig
|
||||||
|
|
||||||
|
from sqlalchemy import engine_from_config, pool
|
||||||
|
from sqlmodel import SQLModel
|
||||||
|
|
||||||
|
from alembic import context
|
||||||
|
from app.config import settings
|
||||||
|
|
||||||
|
# Import all models so they're registered with SQLModel metadata
|
||||||
|
from app.models.User import User # noqa: F401
|
||||||
|
from app.models.Credential import Credential # noqa: F401
|
||||||
|
from app.models.IdentityKey import IdentityKey # noqa: F401
|
||||||
|
from app.models.SignedPrekey import SignedPrekey # noqa: F401
|
||||||
|
from app.models.OneTimePrekey import OneTimePrekey # noqa: F401
|
||||||
|
from app.models.RatchetState import RatchetState # noqa: F401
|
||||||
|
from app.models.SkippedMessageKey import SkippedMessageKey # noqa: F401
|
||||||
|
|
||||||
|
|
||||||
|
config = context.config
|
||||||
|
|
||||||
|
if config.config_file_name is not None:
|
||||||
|
fileConfig(config.config_file_name)
|
||||||
|
|
||||||
|
target_metadata = SQLModel.metadata
|
||||||
|
|
||||||
|
config.set_main_option("sqlalchemy.url", str(settings.DATABASE_URL))
|
||||||
|
|
||||||
|
|
||||||
|
def run_migrations_offline() -> None:
|
||||||
|
"""
|
||||||
|
Run migrations in offline mode
|
||||||
|
"""
|
||||||
|
url = config.get_main_option("sqlalchemy.url")
|
||||||
|
context.configure(
|
||||||
|
url = url,
|
||||||
|
target_metadata = target_metadata,
|
||||||
|
literal_binds = True,
|
||||||
|
dialect_opts = {"paramstyle": "named"},
|
||||||
|
)
|
||||||
|
|
||||||
|
with context.begin_transaction():
|
||||||
|
context.run_migrations()
|
||||||
|
|
||||||
|
|
||||||
|
def run_migrations_online() -> None:
|
||||||
|
"""
|
||||||
|
Run migrations in online mode
|
||||||
|
"""
|
||||||
|
connectable = engine_from_config(
|
||||||
|
config.get_section(config.config_ini_section, {}),
|
||||||
|
prefix = "sqlalchemy.",
|
||||||
|
poolclass = pool.NullPool,
|
||||||
|
)
|
||||||
|
|
||||||
|
with connectable.connect() as connection:
|
||||||
|
context.configure(
|
||||||
|
connection = connection,
|
||||||
|
target_metadata = target_metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
with context.begin_transaction():
|
||||||
|
context.run_migrations()
|
||||||
|
|
||||||
|
|
||||||
|
if context.is_offline_mode():
|
||||||
|
run_migrations_offline()
|
||||||
|
else:
|
||||||
|
run_migrations_online()
|
||||||
|
|
@ -0,0 +1,28 @@
|
||||||
|
"""${message}
|
||||||
|
|
||||||
|
Revision ID: ${up_revision}
|
||||||
|
Revises: ${down_revision | comma,n}
|
||||||
|
Create Date: ${create_date}
|
||||||
|
|
||||||
|
"""
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
from alembic import op
|
||||||
|
import sqlalchemy as sa
|
||||||
|
${imports if imports else ""}
|
||||||
|
|
||||||
|
# revision identifiers, used by Alembic.
|
||||||
|
revision: str = ${repr(up_revision)}
|
||||||
|
down_revision: Union[str, Sequence[str], None] = ${repr(down_revision)}
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
|
||||||
|
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
"""Upgrade schema."""
|
||||||
|
${upgrades if upgrades else "pass"}
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
"""Downgrade schema."""
|
||||||
|
${downgrades if downgrades else "pass"}
|
||||||
|
|
@ -0,0 +1,103 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
WebAuthn passkey registration and login API
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from fastapi import APIRouter, Depends, status
|
||||||
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||||
|
|
||||||
|
from app.models.Base import get_session
|
||||||
|
from app.schemas.auth import (
|
||||||
|
AuthenticationBeginRequest,
|
||||||
|
AuthenticationCompleteRequest,
|
||||||
|
RegistrationBeginRequest,
|
||||||
|
RegistrationCompleteRequest,
|
||||||
|
UserResponse,
|
||||||
|
UserSearchRequest,
|
||||||
|
UserSearchResponse,
|
||||||
|
)
|
||||||
|
from app.services.auth_service import auth_service
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
router = APIRouter(prefix = "/auth", tags = ["authentication"])
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/register/begin", status_code = status.HTTP_200_OK)
|
||||||
|
async def register_begin(
|
||||||
|
request: RegistrationBeginRequest,
|
||||||
|
session: AsyncSession = Depends(get_session),
|
||||||
|
) -> dict[str,
|
||||||
|
Any]:
|
||||||
|
"""
|
||||||
|
Begin WebAuthn passkey registration flow
|
||||||
|
"""
|
||||||
|
return await auth_service.begin_registration(session, request)
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/register/complete", status_code = status.HTTP_201_CREATED)
|
||||||
|
async def register_complete(
|
||||||
|
request: RegistrationCompleteRequest,
|
||||||
|
session: AsyncSession = Depends(get_session),
|
||||||
|
) -> UserResponse:
|
||||||
|
"""
|
||||||
|
Complete WebAuthn passkey registration
|
||||||
|
"""
|
||||||
|
return await auth_service.complete_registration(session, request, request.username)
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/authenticate/begin", status_code = status.HTTP_200_OK)
|
||||||
|
async def authenticate_begin(
|
||||||
|
request: AuthenticationBeginRequest,
|
||||||
|
session: AsyncSession = Depends(get_session),
|
||||||
|
) -> dict[str,
|
||||||
|
Any]:
|
||||||
|
"""
|
||||||
|
Begin WebAuthn passkey authentication flow
|
||||||
|
"""
|
||||||
|
return await auth_service.begin_authentication(session, request)
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/authenticate/complete", status_code = status.HTTP_200_OK)
|
||||||
|
async def authenticate_complete(
|
||||||
|
request: AuthenticationCompleteRequest,
|
||||||
|
session: AsyncSession = Depends(get_session),
|
||||||
|
) -> UserResponse:
|
||||||
|
"""
|
||||||
|
Complete WebAuthn passkey authentication
|
||||||
|
"""
|
||||||
|
return await auth_service.complete_authentication(session, request)
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/users/search", status_code = status.HTTP_200_OK)
|
||||||
|
async def search_users(
|
||||||
|
request: UserSearchRequest,
|
||||||
|
session: AsyncSession = Depends(get_session),
|
||||||
|
) -> UserSearchResponse:
|
||||||
|
"""
|
||||||
|
Search for users by username or display name
|
||||||
|
"""
|
||||||
|
users = await auth_service.search_users(
|
||||||
|
session,
|
||||||
|
request.query,
|
||||||
|
request.limit,
|
||||||
|
)
|
||||||
|
|
||||||
|
return UserSearchResponse(
|
||||||
|
users = [
|
||||||
|
UserResponse(
|
||||||
|
id = str(user.id),
|
||||||
|
username = user.username,
|
||||||
|
display_name = user.display_name,
|
||||||
|
is_active = user.is_active,
|
||||||
|
is_verified = user.is_verified,
|
||||||
|
created_at = user.created_at.isoformat(),
|
||||||
|
)
|
||||||
|
for user in users
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
@ -0,0 +1,89 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Encryption endpoints for X3DH prekey bundles
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from fastapi import APIRouter, Depends, status
|
||||||
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||||
|
|
||||||
|
from app.models.Base import get_session
|
||||||
|
from app.services.prekey_service import prekey_service
|
||||||
|
from app.core.encryption.x3dh_manager import PreKeyBundle
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
router = APIRouter(prefix = "/encryption", tags = ["encryption"])
|
||||||
|
|
||||||
|
|
||||||
|
@router.get(
|
||||||
|
"/prekey-bundle/{user_id}",
|
||||||
|
status_code = status.HTTP_200_OK,
|
||||||
|
response_model = PreKeyBundle
|
||||||
|
)
|
||||||
|
async def get_prekey_bundle(
|
||||||
|
user_id: UUID,
|
||||||
|
session: AsyncSession = Depends(get_session),
|
||||||
|
) -> PreKeyBundle:
|
||||||
|
"""
|
||||||
|
Retrieves prekey bundle for initiating X3DH key exchange with a user
|
||||||
|
"""
|
||||||
|
bundle = await prekey_service.get_prekey_bundle(session, user_id)
|
||||||
|
|
||||||
|
unused_count = await prekey_service.get_unused_opk_count(session, user_id)
|
||||||
|
if unused_count < 20:
|
||||||
|
logger.info("User %s has %s OPKs, replenishing", user_id, unused_count)
|
||||||
|
await prekey_service.replenish_one_time_prekeys(session, user_id, 100)
|
||||||
|
|
||||||
|
return bundle
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/initialize-keys/{user_id}", status_code = status.HTTP_201_CREATED)
|
||||||
|
async def initialize_keys(
|
||||||
|
user_id: UUID,
|
||||||
|
session: AsyncSession = Depends(get_session),
|
||||||
|
) -> dict[str,
|
||||||
|
str]:
|
||||||
|
"""
|
||||||
|
Initializes encryption keys for a user
|
||||||
|
"""
|
||||||
|
await prekey_service.initialize_user_keys(session, user_id)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"status": "success",
|
||||||
|
"message": f"Initialized encryption keys for user {user_id}"
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/rotate-signed-prekey/{user_id}", status_code = status.HTTP_200_OK)
|
||||||
|
async def rotate_signed_prekey(
|
||||||
|
user_id: UUID,
|
||||||
|
session: AsyncSession = Depends(get_session),
|
||||||
|
) -> dict[str,
|
||||||
|
str]:
|
||||||
|
"""
|
||||||
|
Manually rotates signed prekey for a user
|
||||||
|
"""
|
||||||
|
await prekey_service.rotate_signed_prekey(session, user_id)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"status": "success",
|
||||||
|
"message": f"Rotated signed prekey for user {user_id}"
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/opk-count/{user_id}", status_code = status.HTTP_200_OK)
|
||||||
|
async def get_opk_count(
|
||||||
|
user_id: UUID,
|
||||||
|
session: AsyncSession = Depends(get_session),
|
||||||
|
) -> dict[str,
|
||||||
|
int]:
|
||||||
|
"""
|
||||||
|
Returns count of unused one time prekeys for a user
|
||||||
|
"""
|
||||||
|
count = await prekey_service.get_unused_opk_count(session, user_id)
|
||||||
|
|
||||||
|
return {"unused_opks": count}
|
||||||
|
|
@ -0,0 +1,84 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
WebSocket endpoints for real time chat communication
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from fastapi import (
|
||||||
|
APIRouter,
|
||||||
|
WebSocket,
|
||||||
|
WebSocketDisconnect,
|
||||||
|
Query,
|
||||||
|
)
|
||||||
|
from app.core.websocket_manager import connection_manager
|
||||||
|
from app.services.websocket_service import websocket_service
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
router = APIRouter(prefix = "/ws", tags = ["websocket"])
|
||||||
|
|
||||||
|
|
||||||
|
@router.websocket("")
|
||||||
|
async def websocket_endpoint(
|
||||||
|
websocket: WebSocket,
|
||||||
|
user_id: str = Query(...),
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Main WebSocket endpoint for real time messaging
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
user_uuid = UUID(user_id)
|
||||||
|
except ValueError:
|
||||||
|
logger.error("Invalid user_id format: %s", user_id)
|
||||||
|
await websocket.close(code = 1008, reason = "Invalid user ID")
|
||||||
|
return
|
||||||
|
|
||||||
|
connected = await connection_manager.connect(websocket, user_uuid)
|
||||||
|
|
||||||
|
if not connected:
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
data = await websocket.receive_text()
|
||||||
|
|
||||||
|
try:
|
||||||
|
message = json.loads(data)
|
||||||
|
await websocket_service.route_message(
|
||||||
|
websocket,
|
||||||
|
user_uuid,
|
||||||
|
message
|
||||||
|
)
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
logger.error(
|
||||||
|
"Invalid JSON from user %s: %s",
|
||||||
|
user_uuid,
|
||||||
|
data[: 100]
|
||||||
|
)
|
||||||
|
await websocket.send_json(
|
||||||
|
{
|
||||||
|
"type": "error",
|
||||||
|
"error_code": "invalid_json",
|
||||||
|
"error_message": "Invalid JSON format"
|
||||||
|
}
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Error handling message from %s: %s", user_uuid, e)
|
||||||
|
await websocket.send_json(
|
||||||
|
{
|
||||||
|
"type": "error",
|
||||||
|
"error_code": "processing_error",
|
||||||
|
"error_message": str(e)
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
except WebSocketDisconnect:
|
||||||
|
logger.info("WebSocket disconnected for user %s", user_uuid)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("WebSocket error for user %s: %s", user_uuid, e)
|
||||||
|
finally:
|
||||||
|
await connection_manager.disconnect(websocket, user_uuid)
|
||||||
|
|
@ -0,0 +1,222 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
All environment variables and constants are centralized here
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Literal
|
||||||
|
from functools import lru_cache
|
||||||
|
|
||||||
|
from pydantic import (
|
||||||
|
PostgresDsn,
|
||||||
|
RedisDsn,
|
||||||
|
field_validator,
|
||||||
|
ValidationInfo,
|
||||||
|
)
|
||||||
|
from pydantic_settings import (
|
||||||
|
BaseSettings,
|
||||||
|
SettingsConfigDict,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# User field lengths
|
||||||
|
USERNAME_MIN_LENGTH = 3
|
||||||
|
USERNAME_MAX_LENGTH = 50
|
||||||
|
DISPLAY_NAME_MIN_LENGTH = 1
|
||||||
|
DISPLAY_NAME_MAX_LENGTH = 100
|
||||||
|
DEVICE_NAME_MAX_LENGTH = 100
|
||||||
|
PREKEY_MAX_LENGTH = 500
|
||||||
|
|
||||||
|
# User search
|
||||||
|
USER_SEARCH_MIN_LENGTH = 2
|
||||||
|
USER_SEARCH_DEFAULT_LIMIT = 10
|
||||||
|
USER_SEARCH_MAX_LIMIT = 50
|
||||||
|
|
||||||
|
# Credential field lengths
|
||||||
|
CREDENTIAL_ID_MAX_LENGTH = 512
|
||||||
|
PUBLIC_KEY_MAX_LENGTH = 1024
|
||||||
|
AAGUID_MAX_LENGTH = 64
|
||||||
|
ATTESTATION_TYPE_MAX_LENGTH = 50
|
||||||
|
TRANSPORT_MAX_LENGTH = 200
|
||||||
|
|
||||||
|
# Message field lengths
|
||||||
|
MESSAGE_ID_MAX_LENGTH = 64
|
||||||
|
ROOM_ID_MAX_LENGTH = 64
|
||||||
|
ENCRYPTED_CONTENT_MAX_LENGTH = 50000
|
||||||
|
|
||||||
|
# Pagination defaults
|
||||||
|
DEFAULT_MESSAGE_LIMIT = 50
|
||||||
|
MAX_MESSAGE_LIMIT = 200
|
||||||
|
|
||||||
|
# WebSocket message types
|
||||||
|
WS_MESSAGE_TYPE_ENCRYPTED = "encrypted_message"
|
||||||
|
WS_MESSAGE_TYPE_TYPING = "typing"
|
||||||
|
WS_MESSAGE_TYPE_PRESENCE = "presence"
|
||||||
|
WS_MESSAGE_TYPE_RECEIPT = "receipt"
|
||||||
|
WS_MESSAGE_TYPE_ERROR = "error"
|
||||||
|
|
||||||
|
# Encryption key field lengths
|
||||||
|
IDENTITY_KEY_LENGTH = 64
|
||||||
|
SIGNED_PREKEY_LENGTH = 64
|
||||||
|
ONE_TIME_PREKEY_LENGTH = 64
|
||||||
|
SIGNATURE_LENGTH = 128
|
||||||
|
RATCHET_STATE_MAX_LENGTH = 100000
|
||||||
|
|
||||||
|
# Encryption constants
|
||||||
|
X25519_KEY_SIZE = 32
|
||||||
|
ED25519_KEY_SIZE = 32
|
||||||
|
ED25519_SIGNATURE_SIZE = 64
|
||||||
|
AES_GCM_KEY_SIZE = 32
|
||||||
|
AES_GCM_NONCE_SIZE = 12
|
||||||
|
HKDF_OUTPUT_SIZE = 32
|
||||||
|
|
||||||
|
# Double Ratchet limits
|
||||||
|
MAX_SKIP_MESSAGE_KEYS = 1000
|
||||||
|
MAX_CACHED_MESSAGE_KEYS = 2000
|
||||||
|
DEFAULT_ONE_TIME_PREKEY_COUNT = 100
|
||||||
|
SIGNED_PREKEY_ROTATION_HOURS = 48
|
||||||
|
SIGNED_PREKEY_RETENTION_DAYS = 7
|
||||||
|
|
||||||
|
# Server defaults
|
||||||
|
DEFAULT_HOST = "0.0.0.0"
|
||||||
|
DEFAULT_PORT = 8000
|
||||||
|
|
||||||
|
# WebAuthn challenge settings
|
||||||
|
WEBAUTHN_CHALLENGE_TTL_SECONDS = 600
|
||||||
|
WEBAUTHN_CHALLENGE_BYTES = 32
|
||||||
|
|
||||||
|
# Application metadata
|
||||||
|
APP_VERSION = "1.0.0"
|
||||||
|
APP_STATUS = "running"
|
||||||
|
APP_DESCRIPTION = "End to end encrypted P2P chat with Double Ratchet and WebAuthn"
|
||||||
|
|
||||||
|
# Middleware settings
|
||||||
|
GZIP_MINIMUM_SIZE = 1000
|
||||||
|
|
||||||
|
|
||||||
|
class Settings(BaseSettings):
|
||||||
|
"""
|
||||||
|
Application settings with environment variable support
|
||||||
|
"""
|
||||||
|
model_config = SettingsConfigDict(
|
||||||
|
env_file = "../../.env",
|
||||||
|
env_file_encoding = "utf-8",
|
||||||
|
case_sensitive = False,
|
||||||
|
extra = "ignore",
|
||||||
|
)
|
||||||
|
|
||||||
|
ENV: Literal["development", "production", "testing"] = "development"
|
||||||
|
DEBUG: bool = True
|
||||||
|
APP_NAME: str = "encrypted-p2p-chat"
|
||||||
|
SECRET_KEY: str
|
||||||
|
|
||||||
|
POSTGRES_HOST: str = "localhost"
|
||||||
|
POSTGRES_PORT: int = 5432
|
||||||
|
POSTGRES_DB: str = "chat_auth"
|
||||||
|
POSTGRES_USER: str = "chat_user"
|
||||||
|
POSTGRES_PASSWORD: str = ""
|
||||||
|
DATABASE_URL: PostgresDsn | None = None
|
||||||
|
DB_POOL_SIZE: int = 20
|
||||||
|
DB_MAX_OVERFLOW: int = 40
|
||||||
|
|
||||||
|
SURREAL_HOST: str = "localhost"
|
||||||
|
SURREAL_PORT: int = 8000
|
||||||
|
SURREAL_USER: str = "root"
|
||||||
|
SURREAL_PASSWORD: str
|
||||||
|
SURREAL_NAMESPACE: str = "chat"
|
||||||
|
SURREAL_DATABASE: str = "production"
|
||||||
|
SURREAL_URL: str | None = None
|
||||||
|
|
||||||
|
REDIS_HOST: str = "localhost"
|
||||||
|
REDIS_PORT: int = 6379
|
||||||
|
REDIS_PASSWORD: str = ""
|
||||||
|
REDIS_URL: RedisDsn | None = None
|
||||||
|
|
||||||
|
RP_ID: str = "localhost"
|
||||||
|
RP_NAME: str = "Encrypted P2P Chat"
|
||||||
|
RP_ORIGIN: str = "http://localhost:3000"
|
||||||
|
|
||||||
|
CORS_ORIGINS: list[str] = ["http://localhost:3000", "http://localhost:5173"]
|
||||||
|
|
||||||
|
WS_HEARTBEAT_INTERVAL: int = 30
|
||||||
|
WS_MAX_CONNECTIONS_PER_USER: int = 5
|
||||||
|
|
||||||
|
KEY_ROTATION_DAYS: int = 90
|
||||||
|
MAX_SKIPPED_MESSAGE_KEYS: int = 1000
|
||||||
|
|
||||||
|
RATE_LIMIT_MESSAGES_PER_MINUTE: int = 60
|
||||||
|
RATE_LIMIT_AUTH_ATTEMPTS: int = 5
|
||||||
|
|
||||||
|
@field_validator("DATABASE_URL", mode = "before")
|
||||||
|
@classmethod
|
||||||
|
def assemble_db_connection(cls, v: str | None, info: ValidationInfo) -> str:
|
||||||
|
"""
|
||||||
|
Build PostgreSQL connection URL if not provided
|
||||||
|
"""
|
||||||
|
if v:
|
||||||
|
return v
|
||||||
|
data = info.data
|
||||||
|
return (
|
||||||
|
f"postgresql+asyncpg://{data['POSTGRES_USER']}:{data['POSTGRES_PASSWORD']}"
|
||||||
|
f"@{data['POSTGRES_HOST']}:{data['POSTGRES_PORT']}/{data['POSTGRES_DB']}"
|
||||||
|
)
|
||||||
|
|
||||||
|
@field_validator("SURREAL_URL", mode = "before")
|
||||||
|
@classmethod
|
||||||
|
def assemble_surreal_connection(
|
||||||
|
cls,
|
||||||
|
v: str | None,
|
||||||
|
info: ValidationInfo
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Build SurrealDB WebSocket URL if not provided
|
||||||
|
"""
|
||||||
|
if v:
|
||||||
|
return v
|
||||||
|
data = info.data
|
||||||
|
return f"ws://{data['SURREAL_HOST']}:{data['SURREAL_PORT']}"
|
||||||
|
|
||||||
|
@field_validator("REDIS_URL", mode = "before")
|
||||||
|
@classmethod
|
||||||
|
def assemble_redis_connection(
|
||||||
|
cls,
|
||||||
|
v: str | None,
|
||||||
|
info: ValidationInfo
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Build Redis connection URL if not provided
|
||||||
|
"""
|
||||||
|
if v:
|
||||||
|
return v
|
||||||
|
data = info.data
|
||||||
|
password_part = f":{data['REDIS_PASSWORD']}@" if data["REDIS_PASSWORD"
|
||||||
|
] else ""
|
||||||
|
return f"redis://{password_part}{data['REDIS_HOST']}:{data['REDIS_PORT']}"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def is_production(self) -> bool:
|
||||||
|
"""
|
||||||
|
Check if running in production environment
|
||||||
|
"""
|
||||||
|
return self.ENV == "production"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def is_development(self) -> bool:
|
||||||
|
"""
|
||||||
|
Check if running in development environment
|
||||||
|
"""
|
||||||
|
return self.ENV == "development"
|
||||||
|
|
||||||
|
|
||||||
|
@lru_cache
|
||||||
|
def get_settings() -> Settings:
|
||||||
|
"""
|
||||||
|
Get cached settings instance using lru_cache
|
||||||
|
"""
|
||||||
|
return Settings() # type: ignore[call-arg]
|
||||||
|
|
||||||
|
|
||||||
|
settings = get_settings()
|
||||||
|
|
||||||
|
# Export settings fields as module-level constants for imports
|
||||||
|
WS_HEARTBEAT_INTERVAL = settings.WS_HEARTBEAT_INTERVAL
|
||||||
|
WS_MAX_CONNECTIONS_PER_USER = settings.WS_MAX_CONNECTIONS_PER_USER
|
||||||
|
|
@ -0,0 +1,419 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Double Ratchet algorithm implementation for end to end encryption
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
import logging
|
||||||
|
from dataclasses import (
|
||||||
|
field,
|
||||||
|
dataclass,
|
||||||
|
)
|
||||||
|
from cryptography.hazmat.primitives.asymmetric.x25519 import (
|
||||||
|
X25519PrivateKey,
|
||||||
|
X25519PublicKey,
|
||||||
|
)
|
||||||
|
from cryptography.exceptions import InvalidTag
|
||||||
|
from cryptography.hazmat.primitives import hmac
|
||||||
|
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
|
||||||
|
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||||||
|
from cryptography.hazmat.primitives import hashes, serialization
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
AES_GCM_NONCE_SIZE,
|
||||||
|
HKDF_OUTPUT_SIZE,
|
||||||
|
MAX_CACHED_MESSAGE_KEYS,
|
||||||
|
MAX_SKIP_MESSAGE_KEYS,
|
||||||
|
X25519_KEY_SIZE,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DoubleRatchetState:
|
||||||
|
"""
|
||||||
|
Complete state for Double Ratchet algorithm per conversation
|
||||||
|
"""
|
||||||
|
root_key: bytes
|
||||||
|
sending_chain_key: bytes
|
||||||
|
receiving_chain_key: bytes
|
||||||
|
dh_private_key: X25519PrivateKey | None
|
||||||
|
dh_peer_public_key: bytes | None
|
||||||
|
sending_message_number: int = 0
|
||||||
|
receiving_message_number: int = 0
|
||||||
|
previous_sending_chain_length: int = 0
|
||||||
|
skipped_message_keys: dict[tuple[bytes,
|
||||||
|
int],
|
||||||
|
bytes] = field(default_factory = dict)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class EncryptedMessage:
|
||||||
|
"""
|
||||||
|
Encrypted message with header and metadata
|
||||||
|
"""
|
||||||
|
ciphertext: bytes
|
||||||
|
nonce: bytes
|
||||||
|
dh_public_key: bytes
|
||||||
|
message_number: int
|
||||||
|
previous_chain_length: int
|
||||||
|
|
||||||
|
|
||||||
|
class DoubleRatchet:
|
||||||
|
"""
|
||||||
|
Implementation of Signal Protocol Double Ratchet algorithm
|
||||||
|
"""
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
max_skip: int = MAX_SKIP_MESSAGE_KEYS,
|
||||||
|
max_cache: int = MAX_CACHED_MESSAGE_KEYS
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize Double Ratchet with security limits
|
||||||
|
"""
|
||||||
|
self.max_skip = max_skip
|
||||||
|
self.max_cache = max_cache
|
||||||
|
|
||||||
|
def _kdf_rk(self, root_key: bytes, dh_output: bytes) -> tuple[bytes, bytes]:
|
||||||
|
"""
|
||||||
|
Derives new root key and chain key from DH output
|
||||||
|
"""
|
||||||
|
hkdf = HKDF(
|
||||||
|
algorithm = hashes.SHA256(),
|
||||||
|
length = HKDF_OUTPUT_SIZE * 2,
|
||||||
|
salt = root_key,
|
||||||
|
info = b'',
|
||||||
|
)
|
||||||
|
output = hkdf.derive(dh_output)
|
||||||
|
new_root_key = output[: HKDF_OUTPUT_SIZE]
|
||||||
|
new_chain_key = output[HKDF_OUTPUT_SIZE :]
|
||||||
|
|
||||||
|
logger.debug("Derived new root key and chain key")
|
||||||
|
return new_root_key, new_chain_key
|
||||||
|
|
||||||
|
def _kdf_ck(self, chain_key: bytes) -> tuple[bytes, bytes]:
|
||||||
|
"""
|
||||||
|
Derives next chain key and message key from current chain key
|
||||||
|
"""
|
||||||
|
h_chain = hmac.HMAC(chain_key, hashes.SHA256())
|
||||||
|
h_chain.update(b'\x01')
|
||||||
|
next_chain_key = h_chain.finalize()
|
||||||
|
|
||||||
|
h_message = hmac.HMAC(chain_key, hashes.SHA256())
|
||||||
|
h_message.update(b'\x02')
|
||||||
|
message_key = h_message.finalize()
|
||||||
|
|
||||||
|
logger.debug("Derived next chain key and message key")
|
||||||
|
return next_chain_key, message_key
|
||||||
|
|
||||||
|
def _encrypt_with_message_key(
|
||||||
|
self,
|
||||||
|
message_key: bytes,
|
||||||
|
plaintext: bytes,
|
||||||
|
associated_data: bytes
|
||||||
|
) -> tuple[bytes,
|
||||||
|
bytes]:
|
||||||
|
"""
|
||||||
|
Encrypts plaintext using AES-256-GCM with message key
|
||||||
|
"""
|
||||||
|
aesgcm = AESGCM(message_key)
|
||||||
|
nonce = os.urandom(AES_GCM_NONCE_SIZE)
|
||||||
|
ciphertext = aesgcm.encrypt(nonce, plaintext, associated_data)
|
||||||
|
|
||||||
|
logger.debug(
|
||||||
|
"Encrypted %s bytes to %s bytes",
|
||||||
|
len(plaintext),
|
||||||
|
len(ciphertext)
|
||||||
|
)
|
||||||
|
return nonce, ciphertext
|
||||||
|
|
||||||
|
def _decrypt_with_message_key(
|
||||||
|
self,
|
||||||
|
message_key: bytes,
|
||||||
|
nonce: bytes,
|
||||||
|
ciphertext: bytes,
|
||||||
|
associated_data: bytes
|
||||||
|
) -> bytes:
|
||||||
|
"""
|
||||||
|
Decrypts ciphertext using AES-256-GCM with message key
|
||||||
|
"""
|
||||||
|
aesgcm = AESGCM(message_key)
|
||||||
|
try:
|
||||||
|
plaintext = aesgcm.decrypt(nonce, ciphertext, associated_data)
|
||||||
|
logger.debug(
|
||||||
|
"Decrypted %s bytes to %s bytes",
|
||||||
|
len(ciphertext),
|
||||||
|
len(plaintext)
|
||||||
|
)
|
||||||
|
return plaintext
|
||||||
|
except InvalidTag as e:
|
||||||
|
logger.error("Message authentication failed")
|
||||||
|
raise ValueError("Message tampered or corrupted") from e
|
||||||
|
|
||||||
|
def _dh_ratchet_send(self, state: DoubleRatchetState) -> None:
|
||||||
|
"""
|
||||||
|
Performs DH ratchet step when sending
|
||||||
|
"""
|
||||||
|
state.dh_private_key = X25519PrivateKey.generate()
|
||||||
|
|
||||||
|
state.previous_sending_chain_length = state.sending_message_number
|
||||||
|
state.sending_message_number = 0
|
||||||
|
state.receiving_message_number = 0
|
||||||
|
|
||||||
|
if state.dh_peer_public_key:
|
||||||
|
peer_public = X25519PublicKey.from_public_bytes(
|
||||||
|
state.dh_peer_public_key
|
||||||
|
)
|
||||||
|
dh_output = state.dh_private_key.exchange(peer_public)
|
||||||
|
|
||||||
|
state.root_key, state.sending_chain_key = self._kdf_rk(
|
||||||
|
state.root_key,
|
||||||
|
dh_output
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug("DH ratchet step completed (send)")
|
||||||
|
|
||||||
|
def _dh_ratchet_receive(
|
||||||
|
self,
|
||||||
|
state: DoubleRatchetState,
|
||||||
|
peer_public_key: bytes
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Performs DH ratchet step when receiving
|
||||||
|
"""
|
||||||
|
state.previous_sending_chain_length = state.sending_message_number
|
||||||
|
state.sending_message_number = 0
|
||||||
|
state.receiving_message_number = 0
|
||||||
|
state.dh_peer_public_key = peer_public_key
|
||||||
|
|
||||||
|
if state.dh_private_key:
|
||||||
|
peer_public = X25519PublicKey.from_public_bytes(peer_public_key)
|
||||||
|
dh_output = state.dh_private_key.exchange(peer_public)
|
||||||
|
|
||||||
|
state.root_key, state.receiving_chain_key = self._kdf_rk(
|
||||||
|
state.root_key,
|
||||||
|
dh_output
|
||||||
|
)
|
||||||
|
|
||||||
|
state.dh_private_key = X25519PrivateKey.generate()
|
||||||
|
|
||||||
|
if state.dh_peer_public_key:
|
||||||
|
peer_public = X25519PublicKey.from_public_bytes(
|
||||||
|
state.dh_peer_public_key
|
||||||
|
)
|
||||||
|
dh_output = state.dh_private_key.exchange(peer_public)
|
||||||
|
|
||||||
|
state.root_key, state.sending_chain_key = self._kdf_rk(
|
||||||
|
state.root_key,
|
||||||
|
dh_output
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug("DH ratchet step completed (receive)")
|
||||||
|
|
||||||
|
def _store_skipped_message_keys(
|
||||||
|
self,
|
||||||
|
state: DoubleRatchetState,
|
||||||
|
until_message_number: int,
|
||||||
|
dh_public_key: bytes
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Stores skipped message keys for out of order delivery
|
||||||
|
"""
|
||||||
|
num_to_skip = until_message_number - state.receiving_message_number
|
||||||
|
|
||||||
|
if num_to_skip > self.max_skip:
|
||||||
|
raise ValueError(
|
||||||
|
f"Cannot skip {num_to_skip} messages "
|
||||||
|
f"(MAX_SKIP={self.max_skip})"
|
||||||
|
)
|
||||||
|
|
||||||
|
if len(state.skipped_message_keys) + num_to_skip > self.max_cache:
|
||||||
|
logger.warning("Skipped message key cache full, evicting oldest keys")
|
||||||
|
self._evict_oldest_skipped_keys(state, num_to_skip)
|
||||||
|
|
||||||
|
chain_key = state.receiving_chain_key
|
||||||
|
for msg_num in range(state.receiving_message_number,
|
||||||
|
until_message_number):
|
||||||
|
chain_key, message_key = self._kdf_ck(chain_key)
|
||||||
|
state.skipped_message_keys[(dh_public_key, msg_num)] = message_key
|
||||||
|
|
||||||
|
state.receiving_chain_key = chain_key
|
||||||
|
|
||||||
|
logger.debug("Stored %s skipped message keys", num_to_skip)
|
||||||
|
|
||||||
|
def _evict_oldest_skipped_keys(
|
||||||
|
self,
|
||||||
|
state: DoubleRatchetState,
|
||||||
|
count: int
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Evicts oldest skipped message keys to make room
|
||||||
|
"""
|
||||||
|
keys_to_remove = list(state.skipped_message_keys.keys())[: count]
|
||||||
|
for key in keys_to_remove:
|
||||||
|
del state.skipped_message_keys[key]
|
||||||
|
|
||||||
|
logger.debug("Evicted %s skipped keys", len(keys_to_remove))
|
||||||
|
|
||||||
|
def _try_skipped_message_key(
|
||||||
|
self,
|
||||||
|
state: DoubleRatchetState,
|
||||||
|
dh_public_key: bytes,
|
||||||
|
message_number: int
|
||||||
|
) -> bytes | None:
|
||||||
|
"""
|
||||||
|
Attempts to retrieve skipped message key
|
||||||
|
"""
|
||||||
|
key = (dh_public_key, message_number)
|
||||||
|
message_key = state.skipped_message_keys.pop(key, None)
|
||||||
|
|
||||||
|
if message_key:
|
||||||
|
logger.debug(
|
||||||
|
"Retrieved skipped message key for msg %s",
|
||||||
|
message_number
|
||||||
|
)
|
||||||
|
return message_key
|
||||||
|
|
||||||
|
def initialize_sender(
|
||||||
|
self,
|
||||||
|
shared_key: bytes,
|
||||||
|
peer_public_key: bytes
|
||||||
|
) -> DoubleRatchetState:
|
||||||
|
"""
|
||||||
|
Initializes Double Ratchet as sender after X3DH
|
||||||
|
"""
|
||||||
|
dh_private = X25519PrivateKey.generate()
|
||||||
|
peer_public = X25519PublicKey.from_public_bytes(peer_public_key)
|
||||||
|
dh_output = dh_private.exchange(peer_public)
|
||||||
|
|
||||||
|
root_key, sending_chain_key = self._kdf_rk(shared_key, dh_output)
|
||||||
|
|
||||||
|
state = DoubleRatchetState(
|
||||||
|
root_key = root_key,
|
||||||
|
sending_chain_key = sending_chain_key,
|
||||||
|
receiving_chain_key = b'\x00' * HKDF_OUTPUT_SIZE,
|
||||||
|
dh_private_key = dh_private,
|
||||||
|
dh_peer_public_key = peer_public_key
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info("Double Ratchet initialized as sender")
|
||||||
|
return state
|
||||||
|
|
||||||
|
def initialize_receiver(
|
||||||
|
self,
|
||||||
|
shared_key: bytes,
|
||||||
|
own_private_key: X25519PrivateKey
|
||||||
|
) -> DoubleRatchetState:
|
||||||
|
"""
|
||||||
|
Initializes Double Ratchet as receiver after X3DH
|
||||||
|
"""
|
||||||
|
state = DoubleRatchetState(
|
||||||
|
root_key = shared_key,
|
||||||
|
sending_chain_key = b'\x00' * HKDF_OUTPUT_SIZE,
|
||||||
|
receiving_chain_key = b'\x00' * HKDF_OUTPUT_SIZE,
|
||||||
|
dh_private_key = own_private_key,
|
||||||
|
dh_peer_public_key = None
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info("Double Ratchet initialized as receiver")
|
||||||
|
return state
|
||||||
|
|
||||||
|
def encrypt_message(
|
||||||
|
self,
|
||||||
|
state: DoubleRatchetState,
|
||||||
|
plaintext: bytes,
|
||||||
|
associated_data: bytes
|
||||||
|
) -> EncryptedMessage:
|
||||||
|
"""
|
||||||
|
Encrypts message and advances sending ratchet
|
||||||
|
"""
|
||||||
|
state.sending_chain_key, message_key = self._kdf_ck(
|
||||||
|
state.sending_chain_key
|
||||||
|
)
|
||||||
|
|
||||||
|
nonce, ciphertext = self._encrypt_with_message_key(
|
||||||
|
message_key,
|
||||||
|
plaintext,
|
||||||
|
associated_data
|
||||||
|
)
|
||||||
|
|
||||||
|
if state.dh_private_key:
|
||||||
|
dh_public = state.dh_private_key.public_key()
|
||||||
|
dh_public_bytes = dh_public.public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
dh_public_bytes = b'\x00' * X25519_KEY_SIZE
|
||||||
|
|
||||||
|
encrypted_msg = EncryptedMessage(
|
||||||
|
ciphertext = ciphertext,
|
||||||
|
nonce = nonce,
|
||||||
|
dh_public_key = dh_public_bytes,
|
||||||
|
message_number = state.sending_message_number,
|
||||||
|
previous_chain_length = state.previous_sending_chain_length
|
||||||
|
)
|
||||||
|
|
||||||
|
state.sending_message_number += 1
|
||||||
|
|
||||||
|
logger.info("Encrypted message #%s", encrypted_msg.message_number)
|
||||||
|
return encrypted_msg
|
||||||
|
|
||||||
|
def decrypt_message(
|
||||||
|
self,
|
||||||
|
state: DoubleRatchetState,
|
||||||
|
encrypted_msg: EncryptedMessage,
|
||||||
|
associated_data: bytes
|
||||||
|
) -> bytes:
|
||||||
|
"""
|
||||||
|
Decrypts message and advances receiving ratchet
|
||||||
|
"""
|
||||||
|
skipped_key = self._try_skipped_message_key(
|
||||||
|
state,
|
||||||
|
encrypted_msg.dh_public_key,
|
||||||
|
encrypted_msg.message_number
|
||||||
|
)
|
||||||
|
if skipped_key:
|
||||||
|
return self._decrypt_with_message_key(
|
||||||
|
skipped_key,
|
||||||
|
encrypted_msg.nonce,
|
||||||
|
encrypted_msg.ciphertext,
|
||||||
|
associated_data
|
||||||
|
)
|
||||||
|
|
||||||
|
if encrypted_msg.dh_public_key != state.dh_peer_public_key:
|
||||||
|
if state.dh_peer_public_key:
|
||||||
|
self._store_skipped_message_keys(
|
||||||
|
state,
|
||||||
|
encrypted_msg.previous_chain_length,
|
||||||
|
state.dh_peer_public_key
|
||||||
|
)
|
||||||
|
|
||||||
|
self._dh_ratchet_receive(state, encrypted_msg.dh_public_key)
|
||||||
|
|
||||||
|
if encrypted_msg.message_number > state.receiving_message_number:
|
||||||
|
self._store_skipped_message_keys(
|
||||||
|
state,
|
||||||
|
encrypted_msg.message_number,
|
||||||
|
encrypted_msg.dh_public_key
|
||||||
|
)
|
||||||
|
|
||||||
|
state.receiving_chain_key, message_key = self._kdf_ck(
|
||||||
|
state.receiving_chain_key
|
||||||
|
)
|
||||||
|
state.receiving_message_number += 1
|
||||||
|
|
||||||
|
plaintext = self._decrypt_with_message_key(
|
||||||
|
message_key,
|
||||||
|
encrypted_msg.nonce,
|
||||||
|
encrypted_msg.ciphertext,
|
||||||
|
associated_data
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info("Decrypted message #%s", encrypted_msg.message_number)
|
||||||
|
return plaintext
|
||||||
|
|
||||||
|
|
||||||
|
double_ratchet = DoubleRatchet()
|
||||||
|
|
@ -0,0 +1,352 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
X3DH key exchange manager for async initial key agreement
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
from cryptography.hazmat.primitives.asymmetric.ed25519 import (
|
||||||
|
Ed25519PrivateKey,
|
||||||
|
Ed25519PublicKey,
|
||||||
|
)
|
||||||
|
from cryptography.hazmat.primitives import hashes, serialization
|
||||||
|
from cryptography.hazmat.primitives.asymmetric.x25519 import (
|
||||||
|
X25519PrivateKey,
|
||||||
|
X25519PublicKey,
|
||||||
|
)
|
||||||
|
from cryptography.exceptions import InvalidSignature
|
||||||
|
from webauthn.helpers import (
|
||||||
|
bytes_to_base64url,
|
||||||
|
base64url_to_bytes,
|
||||||
|
)
|
||||||
|
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
X25519_KEY_SIZE,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class PreKeyBundle:
|
||||||
|
"""
|
||||||
|
Recipient prekey bundle for X3DH protocol
|
||||||
|
"""
|
||||||
|
identity_key: str
|
||||||
|
signed_prekey: str
|
||||||
|
signed_prekey_signature: str
|
||||||
|
one_time_prekey: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class X3DHResult:
|
||||||
|
"""
|
||||||
|
Result of X3DH key exchange containing shared key and metadata
|
||||||
|
"""
|
||||||
|
shared_key: bytes
|
||||||
|
associated_data: bytes
|
||||||
|
ephemeral_public_key: str
|
||||||
|
used_one_time_prekey: bool
|
||||||
|
|
||||||
|
|
||||||
|
class X3DHManager:
|
||||||
|
"""
|
||||||
|
Manages X3DH key exchange protocol for async initial key agreement
|
||||||
|
"""
|
||||||
|
def generate_identity_keypair_x25519(self) -> tuple[str, str]:
|
||||||
|
"""
|
||||||
|
Generates X25519 identity keypair for DH operations
|
||||||
|
"""
|
||||||
|
private_key = X25519PrivateKey.generate()
|
||||||
|
public_key = private_key.public_key()
|
||||||
|
|
||||||
|
private_bytes = private_key.private_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PrivateFormat.Raw,
|
||||||
|
encryption_algorithm = serialization.NoEncryption()
|
||||||
|
)
|
||||||
|
public_bytes = public_key.public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug(
|
||||||
|
"Generated X25519 identity keypair: %s private, %s public bytes",
|
||||||
|
len(private_bytes),
|
||||||
|
len(public_bytes)
|
||||||
|
)
|
||||||
|
|
||||||
|
return (
|
||||||
|
bytes_to_base64url(private_bytes),
|
||||||
|
bytes_to_base64url(public_bytes)
|
||||||
|
)
|
||||||
|
|
||||||
|
def generate_identity_keypair_ed25519(self) -> tuple[str, str]:
|
||||||
|
"""
|
||||||
|
Generates Ed25519 identity keypair for signing prekeys
|
||||||
|
"""
|
||||||
|
private_key = Ed25519PrivateKey.generate()
|
||||||
|
public_key = private_key.public_key()
|
||||||
|
|
||||||
|
private_bytes = private_key.private_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PrivateFormat.Raw,
|
||||||
|
encryption_algorithm = serialization.NoEncryption()
|
||||||
|
)
|
||||||
|
public_bytes = public_key.public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug(
|
||||||
|
"Generated Ed25519 signing keypair: %s private, %s public bytes",
|
||||||
|
len(private_bytes),
|
||||||
|
len(public_bytes)
|
||||||
|
)
|
||||||
|
|
||||||
|
return (
|
||||||
|
bytes_to_base64url(private_bytes),
|
||||||
|
bytes_to_base64url(public_bytes)
|
||||||
|
)
|
||||||
|
|
||||||
|
def generate_signed_prekey(self,
|
||||||
|
identity_private_key_ed25519: str) -> tuple[str,
|
||||||
|
str,
|
||||||
|
str]:
|
||||||
|
"""
|
||||||
|
Generates signed prekey with signature from Ed25519 identity key
|
||||||
|
"""
|
||||||
|
spk_private = X25519PrivateKey.generate()
|
||||||
|
spk_public = spk_private.public_key()
|
||||||
|
|
||||||
|
spk_private_bytes = spk_private.private_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PrivateFormat.Raw,
|
||||||
|
encryption_algorithm = serialization.NoEncryption()
|
||||||
|
)
|
||||||
|
spk_public_bytes = spk_public.public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
|
||||||
|
identity_private_bytes = base64url_to_bytes(identity_private_key_ed25519)
|
||||||
|
identity_private = Ed25519PrivateKey.from_private_bytes(
|
||||||
|
identity_private_bytes
|
||||||
|
)
|
||||||
|
|
||||||
|
signature = identity_private.sign(spk_public_bytes)
|
||||||
|
|
||||||
|
logger.debug(
|
||||||
|
"Generated signed prekey with %s byte signature",
|
||||||
|
len(signature)
|
||||||
|
)
|
||||||
|
|
||||||
|
return (
|
||||||
|
bytes_to_base64url(spk_private_bytes),
|
||||||
|
bytes_to_base64url(spk_public_bytes),
|
||||||
|
bytes_to_base64url(signature)
|
||||||
|
)
|
||||||
|
|
||||||
|
def generate_one_time_prekey(self) -> tuple[str, str]:
|
||||||
|
"""
|
||||||
|
Generates single-use one-time prekey
|
||||||
|
"""
|
||||||
|
opk_private = X25519PrivateKey.generate()
|
||||||
|
opk_public = opk_private.public_key()
|
||||||
|
|
||||||
|
opk_private_bytes = opk_private.private_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PrivateFormat.Raw,
|
||||||
|
encryption_algorithm = serialization.NoEncryption()
|
||||||
|
)
|
||||||
|
opk_public_bytes = opk_public.public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
|
||||||
|
return (
|
||||||
|
bytes_to_base64url(opk_private_bytes),
|
||||||
|
bytes_to_base64url(opk_public_bytes)
|
||||||
|
)
|
||||||
|
|
||||||
|
def verify_signed_prekey(
|
||||||
|
self,
|
||||||
|
signed_prekey_public: str,
|
||||||
|
signature: str,
|
||||||
|
identity_public_key_ed25519: str
|
||||||
|
) -> bool:
|
||||||
|
"""
|
||||||
|
Verifies signed prekey signature using Ed25519 identity key
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
spk_public_bytes = base64url_to_bytes(signed_prekey_public)
|
||||||
|
signature_bytes = base64url_to_bytes(signature)
|
||||||
|
identity_public_bytes = base64url_to_bytes(
|
||||||
|
identity_public_key_ed25519
|
||||||
|
)
|
||||||
|
|
||||||
|
identity_public = Ed25519PublicKey.from_public_bytes(
|
||||||
|
identity_public_bytes
|
||||||
|
)
|
||||||
|
|
||||||
|
identity_public.verify(signature_bytes, spk_public_bytes)
|
||||||
|
|
||||||
|
logger.debug("Signed prekey signature verified successfully")
|
||||||
|
return True
|
||||||
|
|
||||||
|
except InvalidSignature:
|
||||||
|
logger.warning("Signed prekey signature verification failed")
|
||||||
|
return False
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Error verifying signed prekey: %s", e)
|
||||||
|
return False
|
||||||
|
|
||||||
|
def perform_x3dh_sender(
|
||||||
|
self,
|
||||||
|
alice_identity_private_x25519: str,
|
||||||
|
bob_bundle: PreKeyBundle,
|
||||||
|
bob_identity_public_ed25519: str
|
||||||
|
) -> X3DHResult:
|
||||||
|
"""
|
||||||
|
Performs X3DH key exchange from sender side
|
||||||
|
"""
|
||||||
|
if not self.verify_signed_prekey(bob_bundle.signed_prekey,
|
||||||
|
bob_bundle.signed_prekey_signature,
|
||||||
|
bob_identity_public_ed25519):
|
||||||
|
raise ValueError("Invalid signed prekey signature")
|
||||||
|
|
||||||
|
alice_ik_private_bytes = base64url_to_bytes(alice_identity_private_x25519)
|
||||||
|
alice_ik_private = X25519PrivateKey.from_private_bytes(
|
||||||
|
alice_ik_private_bytes
|
||||||
|
)
|
||||||
|
|
||||||
|
alice_ek_private = X25519PrivateKey.generate()
|
||||||
|
alice_ek_public = alice_ek_private.public_key()
|
||||||
|
|
||||||
|
alice_ek_public_bytes = alice_ek_public.public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
|
||||||
|
bob_ik_public_bytes = base64url_to_bytes(bob_bundle.identity_key)
|
||||||
|
bob_ik_public = X25519PublicKey.from_public_bytes(bob_ik_public_bytes)
|
||||||
|
|
||||||
|
bob_spk_public_bytes = base64url_to_bytes(bob_bundle.signed_prekey)
|
||||||
|
bob_spk_public = X25519PublicKey.from_public_bytes(bob_spk_public_bytes)
|
||||||
|
|
||||||
|
dh1 = alice_ik_private.exchange(bob_spk_public)
|
||||||
|
dh2 = alice_ek_private.exchange(bob_ik_public)
|
||||||
|
dh3 = alice_ek_private.exchange(bob_spk_public)
|
||||||
|
|
||||||
|
used_one_time_prekey = False
|
||||||
|
if bob_bundle.one_time_prekey:
|
||||||
|
bob_opk_public_bytes = base64url_to_bytes(bob_bundle.one_time_prekey)
|
||||||
|
bob_opk_public = X25519PublicKey.from_public_bytes(
|
||||||
|
bob_opk_public_bytes
|
||||||
|
)
|
||||||
|
dh4 = alice_ek_private.exchange(bob_opk_public)
|
||||||
|
key_material = dh1 + dh2 + dh3 + dh4
|
||||||
|
used_one_time_prekey = True
|
||||||
|
else:
|
||||||
|
key_material = dh1 + dh2 + dh3
|
||||||
|
|
||||||
|
f = b'\xff' * X25519_KEY_SIZE
|
||||||
|
hkdf = HKDF(
|
||||||
|
algorithm = hashes.SHA256(),
|
||||||
|
length = X25519_KEY_SIZE,
|
||||||
|
salt = b'\x00' * X25519_KEY_SIZE,
|
||||||
|
info = b'X3DH',
|
||||||
|
)
|
||||||
|
shared_key = hkdf.derive(f + key_material)
|
||||||
|
|
||||||
|
alice_ik_public = alice_ik_private.public_key()
|
||||||
|
alice_ik_public_bytes = alice_ik_public.public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
|
||||||
|
associated_data = alice_ik_public_bytes + bob_ik_public_bytes
|
||||||
|
|
||||||
|
logger.info("X3DH sender completed: OPK used=%s", used_one_time_prekey)
|
||||||
|
|
||||||
|
return X3DHResult(
|
||||||
|
shared_key = shared_key,
|
||||||
|
associated_data = associated_data,
|
||||||
|
ephemeral_public_key = bytes_to_base64url(alice_ek_public_bytes),
|
||||||
|
used_one_time_prekey = used_one_time_prekey
|
||||||
|
)
|
||||||
|
|
||||||
|
def perform_x3dh_receiver(
|
||||||
|
self,
|
||||||
|
bob_identity_private_x25519: str,
|
||||||
|
bob_signed_prekey_private: str,
|
||||||
|
bob_one_time_prekey_private: str | None,
|
||||||
|
alice_identity_public_x25519: str,
|
||||||
|
alice_ephemeral_public: str
|
||||||
|
) -> X3DHResult:
|
||||||
|
"""
|
||||||
|
Performs X3DH key exchange from receiver side
|
||||||
|
"""
|
||||||
|
bob_ik_private_bytes = base64url_to_bytes(bob_identity_private_x25519)
|
||||||
|
bob_ik_private = X25519PrivateKey.from_private_bytes(bob_ik_private_bytes)
|
||||||
|
|
||||||
|
bob_spk_private_bytes = base64url_to_bytes(bob_signed_prekey_private)
|
||||||
|
bob_spk_private = X25519PrivateKey.from_private_bytes(
|
||||||
|
bob_spk_private_bytes
|
||||||
|
)
|
||||||
|
|
||||||
|
alice_ik_public_bytes = base64url_to_bytes(alice_identity_public_x25519)
|
||||||
|
alice_ik_public = X25519PublicKey.from_public_bytes(alice_ik_public_bytes)
|
||||||
|
|
||||||
|
alice_ek_public_bytes = base64url_to_bytes(alice_ephemeral_public)
|
||||||
|
alice_ek_public = X25519PublicKey.from_public_bytes(alice_ek_public_bytes)
|
||||||
|
|
||||||
|
dh1 = bob_spk_private.exchange(alice_ik_public)
|
||||||
|
dh2 = bob_ik_private.exchange(alice_ek_public)
|
||||||
|
dh3 = bob_spk_private.exchange(alice_ek_public)
|
||||||
|
|
||||||
|
used_one_time_prekey = False
|
||||||
|
if bob_one_time_prekey_private:
|
||||||
|
bob_opk_private_bytes = base64url_to_bytes(
|
||||||
|
bob_one_time_prekey_private
|
||||||
|
)
|
||||||
|
bob_opk_private = X25519PrivateKey.from_private_bytes(
|
||||||
|
bob_opk_private_bytes
|
||||||
|
)
|
||||||
|
dh4 = bob_opk_private.exchange(alice_ek_public)
|
||||||
|
key_material = dh1 + dh2 + dh3 + dh4
|
||||||
|
used_one_time_prekey = True
|
||||||
|
else:
|
||||||
|
key_material = dh1 + dh2 + dh3
|
||||||
|
|
||||||
|
f = b'\xff' * X25519_KEY_SIZE
|
||||||
|
hkdf = HKDF(
|
||||||
|
algorithm = hashes.SHA256(),
|
||||||
|
length = X25519_KEY_SIZE,
|
||||||
|
salt = b'\x00' * X25519_KEY_SIZE,
|
||||||
|
info = b'X3DH',
|
||||||
|
)
|
||||||
|
shared_key = hkdf.derive(f + key_material)
|
||||||
|
|
||||||
|
bob_ik_public = bob_ik_private.public_key()
|
||||||
|
bob_ik_public_bytes = bob_ik_public.public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
|
||||||
|
associated_data = alice_ik_public_bytes + bob_ik_public_bytes
|
||||||
|
|
||||||
|
logger.info("X3DH receiver completed: OPK used=%s", used_one_time_prekey)
|
||||||
|
|
||||||
|
return X3DHResult(
|
||||||
|
shared_key = shared_key,
|
||||||
|
associated_data = associated_data,
|
||||||
|
ephemeral_public_key = alice_ephemeral_public,
|
||||||
|
used_one_time_prekey = used_one_time_prekey
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
x3dh_manager = X3DHManager()
|
||||||
|
|
@ -0,0 +1,35 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Application enums for type safety
|
||||||
|
"""
|
||||||
|
|
||||||
|
from enum import Enum
|
||||||
|
|
||||||
|
|
||||||
|
class MessageStatus(str, Enum):
|
||||||
|
"""
|
||||||
|
Message delivery status
|
||||||
|
"""
|
||||||
|
SENDING = "sending"
|
||||||
|
SENT = "sent"
|
||||||
|
DELIVERED = "delivered"
|
||||||
|
READ = "read"
|
||||||
|
FAILED = "failed"
|
||||||
|
|
||||||
|
|
||||||
|
class PresenceStatus(str, Enum):
|
||||||
|
"""
|
||||||
|
User presence status
|
||||||
|
"""
|
||||||
|
ONLINE = "online"
|
||||||
|
AWAY = "away"
|
||||||
|
OFFLINE = "offline"
|
||||||
|
|
||||||
|
|
||||||
|
class RoomType(str, Enum):
|
||||||
|
"""
|
||||||
|
Chat room types
|
||||||
|
"""
|
||||||
|
DIRECT = "direct"
|
||||||
|
GROUP = "group"
|
||||||
|
EPHEMERAL = "ephemeral"
|
||||||
|
|
@ -0,0 +1,246 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Global exception handlers for FastAPI application
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from fastapi import FastAPI, Request, status
|
||||||
|
from fastapi.responses import JSONResponse
|
||||||
|
|
||||||
|
from app.core.exceptions import (
|
||||||
|
AuthenticationError,
|
||||||
|
ChallengeExpiredError,
|
||||||
|
CredentialNotFoundError,
|
||||||
|
CredentialVerificationError,
|
||||||
|
DatabaseError,
|
||||||
|
DecryptionError,
|
||||||
|
EncryptionError,
|
||||||
|
InvalidDataError,
|
||||||
|
KeyExchangeError,
|
||||||
|
RatchetStateNotFoundError,
|
||||||
|
UserExistsError,
|
||||||
|
UserInactiveError,
|
||||||
|
UserNotFoundError,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
async def user_exists_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: UserExistsError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle UserExistsError exceptions
|
||||||
|
"""
|
||||||
|
logger.warning("User exists error on %s: %s", request.url, exc.message)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_409_CONFLICT,
|
||||||
|
content = {"detail": exc.message},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def user_not_found_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: UserNotFoundError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle UserNotFoundError exceptions
|
||||||
|
"""
|
||||||
|
logger.warning("User not found on %s: %s", request.url, exc.message)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_404_NOT_FOUND,
|
||||||
|
content = {"detail": exc.message},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def user_inactive_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: UserInactiveError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle UserInactiveError exceptions
|
||||||
|
"""
|
||||||
|
logger.warning(
|
||||||
|
"Inactive user access attempt on %s: %s",
|
||||||
|
request.url,
|
||||||
|
exc.message
|
||||||
|
)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_403_FORBIDDEN,
|
||||||
|
content = {"detail": exc.message},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def credential_not_found_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: CredentialNotFoundError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle CredentialNotFoundError exceptions
|
||||||
|
"""
|
||||||
|
logger.warning("Credential not found on %s: %s", request.url, exc.message)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_404_NOT_FOUND,
|
||||||
|
content = {"detail": exc.message},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def credential_verification_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: CredentialVerificationError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle CredentialVerificationError exceptions
|
||||||
|
"""
|
||||||
|
logger.error(
|
||||||
|
"Credential verification failed on %s: %s",
|
||||||
|
request.url,
|
||||||
|
exc.message
|
||||||
|
)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_401_UNAUTHORIZED,
|
||||||
|
content = {"detail": exc.message},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def challenge_expired_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: ChallengeExpiredError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle ChallengeExpiredError exceptions
|
||||||
|
"""
|
||||||
|
logger.warning("Challenge expired on %s: %s", request.url, exc.message)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_400_BAD_REQUEST,
|
||||||
|
content = {"detail": exc.message},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def database_error_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: DatabaseError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle DatabaseError exceptions
|
||||||
|
"""
|
||||||
|
logger.error("Database error on %s: %s", request.url, exc.message)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||||
|
content = {"detail": "Internal server error"},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def authentication_error_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: AuthenticationError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle AuthenticationError exceptions
|
||||||
|
"""
|
||||||
|
logger.warning("Authentication error on %s: %s", request.url, exc.message)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_401_UNAUTHORIZED,
|
||||||
|
content = {"detail": exc.message},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def invalid_data_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: InvalidDataError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle InvalidDataError exceptions
|
||||||
|
"""
|
||||||
|
logger.warning("Invalid data on %s: %s", request.url, exc.message)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_400_BAD_REQUEST,
|
||||||
|
content = {"detail": exc.message},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def encryption_error_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: EncryptionError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle EncryptionError exceptions
|
||||||
|
"""
|
||||||
|
logger.error("Encryption error on %s: %s", request.url, exc.message)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||||
|
content = {"detail": "Encryption failed"},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def decryption_error_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: DecryptionError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle DecryptionError exceptions
|
||||||
|
"""
|
||||||
|
logger.error("Decryption error on %s: %s", request.url, exc.message)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||||
|
content = {"detail": "Decryption failed"},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def ratchet_state_not_found_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: RatchetStateNotFoundError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle RatchetStateNotFoundError exceptions
|
||||||
|
"""
|
||||||
|
logger.warning("Ratchet state not found on %s: %s", request.url, exc.message)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_404_NOT_FOUND,
|
||||||
|
content = {"detail": exc.message},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def key_exchange_error_handler(
|
||||||
|
request: Request,
|
||||||
|
exc: KeyExchangeError
|
||||||
|
) -> JSONResponse:
|
||||||
|
"""
|
||||||
|
Handle KeyExchangeError exceptions
|
||||||
|
"""
|
||||||
|
logger.error("Key exchange error on %s: %s", request.url, exc.message)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code = status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||||
|
content = {"detail": "Key exchange failed"},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def register_exception_handlers(app: FastAPI) -> None:
|
||||||
|
"""
|
||||||
|
Register all custom exception handlers with FastAPI app
|
||||||
|
"""
|
||||||
|
app.add_exception_handler(UserExistsError, user_exists_handler)
|
||||||
|
app.add_exception_handler(UserNotFoundError, user_not_found_handler)
|
||||||
|
app.add_exception_handler(UserInactiveError, user_inactive_handler)
|
||||||
|
app.add_exception_handler(
|
||||||
|
CredentialNotFoundError,
|
||||||
|
credential_not_found_handler
|
||||||
|
)
|
||||||
|
app.add_exception_handler(
|
||||||
|
CredentialVerificationError,
|
||||||
|
credential_verification_handler
|
||||||
|
)
|
||||||
|
app.add_exception_handler(ChallengeExpiredError, challenge_expired_handler)
|
||||||
|
app.add_exception_handler(DatabaseError, database_error_handler)
|
||||||
|
app.add_exception_handler(AuthenticationError, authentication_error_handler)
|
||||||
|
app.add_exception_handler(InvalidDataError, invalid_data_handler)
|
||||||
|
app.add_exception_handler(EncryptionError, encryption_error_handler)
|
||||||
|
app.add_exception_handler(DecryptionError, decryption_error_handler)
|
||||||
|
app.add_exception_handler(
|
||||||
|
RatchetStateNotFoundError,
|
||||||
|
ratchet_state_not_found_handler
|
||||||
|
)
|
||||||
|
app.add_exception_handler(KeyExchangeError, key_exchange_error_handler)
|
||||||
|
|
@ -0,0 +1,94 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Custom application exceptions for clean error handling
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class AppException(Exception):
|
||||||
|
"""
|
||||||
|
Base application exception
|
||||||
|
"""
|
||||||
|
def __init__(self, message: str) -> None:
|
||||||
|
"""
|
||||||
|
Initialize exception with message
|
||||||
|
"""
|
||||||
|
self.message = message
|
||||||
|
super().__init__(self.message)
|
||||||
|
|
||||||
|
|
||||||
|
class UserExistsError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when attempting to create a user that already exists
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class UserNotFoundError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when user cannot be found
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class UserInactiveError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when user account is inactive
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class CredentialNotFoundError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when credential cannot be found
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class CredentialVerificationError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when credential verification fails
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class ChallengeExpiredError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when WebAuthn challenge has expired or not found
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class DatabaseError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when database operation fails
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class AuthenticationError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when authentication fails
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class InvalidDataError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when input data is invalid
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class EncryptionError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when message encryption fails
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class DecryptionError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when message decryption fails
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class RatchetStateNotFoundError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when ratchet state not found for conversation
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class KeyExchangeError(AppException):
|
||||||
|
"""
|
||||||
|
Raised when X3DH key exchange fails
|
||||||
|
"""
|
||||||
|
|
@ -0,0 +1,210 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
WebAuthn passkey manager using py_webauthn library
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import secrets
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from webauthn.helpers import (
|
||||||
|
bytes_to_base64url,
|
||||||
|
options_to_json_dict,
|
||||||
|
)
|
||||||
|
from webauthn import (
|
||||||
|
generate_authentication_options,
|
||||||
|
generate_registration_options,
|
||||||
|
verify_authentication_response,
|
||||||
|
verify_registration_response,
|
||||||
|
)
|
||||||
|
from webauthn.helpers.structs import (
|
||||||
|
AttestationConveyancePreference,
|
||||||
|
AuthenticatorSelectionCriteria,
|
||||||
|
PublicKeyCredentialDescriptor,
|
||||||
|
ResidentKeyRequirement,
|
||||||
|
UserVerificationRequirement,
|
||||||
|
)
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
settings,
|
||||||
|
WEBAUTHN_CHALLENGE_BYTES,
|
||||||
|
)
|
||||||
|
from app.schemas.auth import (
|
||||||
|
AuthenticationOptionsResponse,
|
||||||
|
RegistrationOptionsResponse,
|
||||||
|
VerifiedAuthentication,
|
||||||
|
VerifiedRegistration,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class PasskeyManager:
|
||||||
|
"""
|
||||||
|
WebAuthn passkey manager for registration and authentication
|
||||||
|
"""
|
||||||
|
def __init__(self) -> None:
|
||||||
|
"""
|
||||||
|
Initialize passkey manager with RP configuration
|
||||||
|
"""
|
||||||
|
self.rp_id = settings.RP_ID
|
||||||
|
self.rp_name = settings.RP_NAME
|
||||||
|
self.rp_origin = settings.RP_ORIGIN
|
||||||
|
|
||||||
|
def generate_registration_options(
|
||||||
|
self,
|
||||||
|
user_id: bytes,
|
||||||
|
username: str,
|
||||||
|
display_name: str,
|
||||||
|
exclude_credentials: list[bytes] | None = None,
|
||||||
|
) -> RegistrationOptionsResponse:
|
||||||
|
"""
|
||||||
|
Generate WebAuthn registration options for passkey creation
|
||||||
|
"""
|
||||||
|
challenge = secrets.token_bytes(WEBAUTHN_CHALLENGE_BYTES)
|
||||||
|
|
||||||
|
exclude_creds = []
|
||||||
|
if exclude_credentials:
|
||||||
|
exclude_creds = [
|
||||||
|
PublicKeyCredentialDescriptor(id = cred_id)
|
||||||
|
for cred_id in exclude_credentials
|
||||||
|
]
|
||||||
|
|
||||||
|
options = generate_registration_options(
|
||||||
|
rp_id = self.rp_id,
|
||||||
|
rp_name = self.rp_name,
|
||||||
|
user_id = user_id,
|
||||||
|
user_name = username,
|
||||||
|
user_display_name = display_name,
|
||||||
|
challenge = challenge,
|
||||||
|
attestation = AttestationConveyancePreference.NONE,
|
||||||
|
authenticator_selection = AuthenticatorSelectionCriteria(
|
||||||
|
resident_key = ResidentKeyRequirement.REQUIRED,
|
||||||
|
user_verification = UserVerificationRequirement.PREFERRED,
|
||||||
|
),
|
||||||
|
exclude_credentials = exclude_creds,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug("Generated registration options for user %s", username)
|
||||||
|
|
||||||
|
return RegistrationOptionsResponse(
|
||||||
|
options = options_to_json_dict(options),
|
||||||
|
challenge = challenge,
|
||||||
|
)
|
||||||
|
|
||||||
|
def verify_registration(
|
||||||
|
self,
|
||||||
|
credential: dict[str,
|
||||||
|
Any],
|
||||||
|
expected_challenge: bytes,
|
||||||
|
) -> VerifiedRegistration:
|
||||||
|
"""
|
||||||
|
Verify WebAuthn registration response
|
||||||
|
"""
|
||||||
|
verified_registration = verify_registration_response(
|
||||||
|
credential = credential,
|
||||||
|
expected_challenge = expected_challenge,
|
||||||
|
expected_rp_id = self.rp_id,
|
||||||
|
expected_origin = self.rp_origin,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Verified registration for credential %s...",
|
||||||
|
bytes_to_base64url(verified_registration.credential_id)[: 16]
|
||||||
|
)
|
||||||
|
|
||||||
|
return VerifiedRegistration(
|
||||||
|
credential_id = verified_registration.credential_id,
|
||||||
|
credential_public_key = verified_registration.credential_public_key,
|
||||||
|
sign_count = verified_registration.sign_count,
|
||||||
|
aaguid = verified_registration.aaguid,
|
||||||
|
attestation_object = verified_registration.attestation_object,
|
||||||
|
credential_type = verified_registration.credential_type,
|
||||||
|
user_verified = verified_registration.user_verified,
|
||||||
|
attestation_format = verified_registration.fmt,
|
||||||
|
credential_device_type = verified_registration.credential_device_type,
|
||||||
|
credential_backed_up = verified_registration.credential_backed_up,
|
||||||
|
backup_eligible = verified_registration.credential_backed_up,
|
||||||
|
backup_state = verified_registration.credential_backed_up,
|
||||||
|
)
|
||||||
|
|
||||||
|
def generate_authentication_options(
|
||||||
|
self,
|
||||||
|
allow_credentials: list[bytes] | None = None,
|
||||||
|
) -> AuthenticationOptionsResponse:
|
||||||
|
"""
|
||||||
|
Generate WebAuthn authentication options for passkey verification
|
||||||
|
"""
|
||||||
|
challenge = secrets.token_bytes(WEBAUTHN_CHALLENGE_BYTES)
|
||||||
|
|
||||||
|
allow_creds = None
|
||||||
|
if allow_credentials:
|
||||||
|
allow_creds = [
|
||||||
|
PublicKeyCredentialDescriptor(id = cred_id)
|
||||||
|
for cred_id in allow_credentials
|
||||||
|
]
|
||||||
|
|
||||||
|
options = generate_authentication_options(
|
||||||
|
rp_id = self.rp_id,
|
||||||
|
challenge = challenge,
|
||||||
|
allow_credentials = allow_creds,
|
||||||
|
user_verification = UserVerificationRequirement.PREFERRED,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug("Generated authentication options")
|
||||||
|
|
||||||
|
return AuthenticationOptionsResponse(
|
||||||
|
options = options_to_json_dict(options),
|
||||||
|
challenge = challenge,
|
||||||
|
)
|
||||||
|
|
||||||
|
def verify_authentication(
|
||||||
|
self,
|
||||||
|
credential: dict[str,
|
||||||
|
Any],
|
||||||
|
expected_challenge: bytes,
|
||||||
|
credential_public_key: bytes,
|
||||||
|
credential_current_sign_count: int,
|
||||||
|
) -> VerifiedAuthentication:
|
||||||
|
"""
|
||||||
|
Verify WebAuthn authentication response and check signature counter
|
||||||
|
"""
|
||||||
|
verified_authentication = verify_authentication_response(
|
||||||
|
credential = credential,
|
||||||
|
expected_challenge = expected_challenge,
|
||||||
|
expected_rp_id = self.rp_id,
|
||||||
|
expected_origin = self.rp_origin,
|
||||||
|
credential_public_key = credential_public_key,
|
||||||
|
credential_current_sign_count = credential_current_sign_count,
|
||||||
|
)
|
||||||
|
|
||||||
|
new_sign_count = verified_authentication.new_sign_count
|
||||||
|
|
||||||
|
if (credential_current_sign_count != 0 and new_sign_count != 0
|
||||||
|
and new_sign_count <= credential_current_sign_count):
|
||||||
|
logger.error(
|
||||||
|
"Signature counter did not increase: current=%s, new=%s. Possible cloned authenticator detected!",
|
||||||
|
credential_current_sign_count,
|
||||||
|
new_sign_count
|
||||||
|
)
|
||||||
|
raise ValueError(
|
||||||
|
"Signature counter anomaly detected - potential cloned authenticator"
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Verified authentication with counter %s -> %s",
|
||||||
|
credential_current_sign_count,
|
||||||
|
new_sign_count
|
||||||
|
)
|
||||||
|
|
||||||
|
return VerifiedAuthentication(
|
||||||
|
new_sign_count = new_sign_count,
|
||||||
|
credential_id = verified_authentication.credential_id,
|
||||||
|
user_verified = verified_authentication.user_verified,
|
||||||
|
backup_state = verified_authentication.credential_backed_up,
|
||||||
|
backup_eligible = verified_authentication.credential_backup_eligible,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
passkey_manager = PasskeyManager()
|
||||||
|
|
@ -0,0 +1,174 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Redis manager for WebAuthn challenge storage with TTL
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
|
||||||
|
import redis.asyncio as redis
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
settings,
|
||||||
|
WEBAUTHN_CHALLENGE_TTL_SECONDS,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class RedisManager:
|
||||||
|
"""
|
||||||
|
Redis manager for challenge storage with automatic expiration
|
||||||
|
"""
|
||||||
|
def __init__(self) -> None:
|
||||||
|
"""
|
||||||
|
Initialize Redis manager with connection pool
|
||||||
|
"""
|
||||||
|
self.pool: redis.ConnectionPool | None = None
|
||||||
|
self.client: redis.Redis | None = None
|
||||||
|
|
||||||
|
async def connect(self) -> None:
|
||||||
|
"""
|
||||||
|
Establish Redis connection with connection pooling
|
||||||
|
"""
|
||||||
|
if self.pool is not None:
|
||||||
|
return
|
||||||
|
|
||||||
|
self.pool = redis.ConnectionPool.from_url(
|
||||||
|
str(settings.REDIS_URL),
|
||||||
|
max_connections = 50,
|
||||||
|
decode_responses = False,
|
||||||
|
)
|
||||||
|
self.client = redis.Redis(connection_pool = self.pool)
|
||||||
|
|
||||||
|
await self.client.ping()
|
||||||
|
logger.info("Connected to Redis at %s", settings.REDIS_URL)
|
||||||
|
|
||||||
|
async def disconnect(self) -> None:
|
||||||
|
"""
|
||||||
|
Close Redis connection
|
||||||
|
"""
|
||||||
|
if self.client:
|
||||||
|
await self.client.aclose()
|
||||||
|
if self.pool:
|
||||||
|
await self.pool.aclose()
|
||||||
|
logger.info("Disconnected from Redis")
|
||||||
|
|
||||||
|
async def set_registration_challenge(
|
||||||
|
self,
|
||||||
|
user_id: str,
|
||||||
|
challenge: bytes,
|
||||||
|
ttl: int = WEBAUTHN_CHALLENGE_TTL_SECONDS,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Store registration challenge with TTL
|
||||||
|
"""
|
||||||
|
if not self.client:
|
||||||
|
raise RuntimeError("Redis client not connected")
|
||||||
|
|
||||||
|
key = f"webauthn:reg_challenge:{user_id}"
|
||||||
|
await self.client.setex(key, ttl, challenge.hex())
|
||||||
|
logger.debug(
|
||||||
|
"Stored registration challenge for user %s with %ss TTL",
|
||||||
|
user_id,
|
||||||
|
ttl
|
||||||
|
)
|
||||||
|
|
||||||
|
async def get_registration_challenge(self, user_id: str) -> bytes | None:
|
||||||
|
"""
|
||||||
|
Retrieve and delete registration challenge (one-time use)
|
||||||
|
"""
|
||||||
|
if not self.client:
|
||||||
|
raise RuntimeError("Redis client not connected")
|
||||||
|
|
||||||
|
key = f"webauthn:reg_challenge:{user_id}"
|
||||||
|
|
||||||
|
async with self.client.pipeline() as pipe:
|
||||||
|
await pipe.get(key)
|
||||||
|
await pipe.delete(key)
|
||||||
|
results = await pipe.execute()
|
||||||
|
|
||||||
|
challenge_hex = results[0]
|
||||||
|
if challenge_hex is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
return bytes.fromhex(challenge_hex.decode())
|
||||||
|
|
||||||
|
async def set_authentication_challenge(
|
||||||
|
self,
|
||||||
|
user_id: str,
|
||||||
|
challenge: bytes,
|
||||||
|
ttl: int = WEBAUTHN_CHALLENGE_TTL_SECONDS,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Store authentication challenge with TTL
|
||||||
|
"""
|
||||||
|
if not self.client:
|
||||||
|
raise RuntimeError("Redis client not connected")
|
||||||
|
|
||||||
|
key = f"webauthn:auth_challenge:{user_id}"
|
||||||
|
await self.client.setex(key, ttl, challenge.hex())
|
||||||
|
logger.debug(
|
||||||
|
"Stored authentication challenge for user %s with %ss TTL",
|
||||||
|
user_id,
|
||||||
|
ttl
|
||||||
|
)
|
||||||
|
|
||||||
|
async def get_authentication_challenge(self, user_id: str) -> bytes | None:
|
||||||
|
"""
|
||||||
|
Retrieve and delete authentication challenge (one-time use)
|
||||||
|
"""
|
||||||
|
if not self.client:
|
||||||
|
raise RuntimeError("Redis client not connected")
|
||||||
|
|
||||||
|
key = f"webauthn:auth_challenge:{user_id}"
|
||||||
|
|
||||||
|
async with self.client.pipeline() as pipe:
|
||||||
|
await pipe.get(key)
|
||||||
|
await pipe.delete(key)
|
||||||
|
results = await pipe.execute()
|
||||||
|
|
||||||
|
challenge_hex = results[0]
|
||||||
|
if challenge_hex is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
return bytes.fromhex(challenge_hex.decode())
|
||||||
|
|
||||||
|
async def set_value(
|
||||||
|
self,
|
||||||
|
key: str,
|
||||||
|
value: str,
|
||||||
|
ttl: int | None = None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Generic set with optional TTL
|
||||||
|
"""
|
||||||
|
if not self.client:
|
||||||
|
raise RuntimeError("Redis client not connected")
|
||||||
|
|
||||||
|
if ttl:
|
||||||
|
await self.client.setex(key, ttl, value)
|
||||||
|
else:
|
||||||
|
await self.client.set(key, value)
|
||||||
|
|
||||||
|
async def get_value(self, key: str) -> str | None:
|
||||||
|
"""
|
||||||
|
Generic get
|
||||||
|
"""
|
||||||
|
if not self.client:
|
||||||
|
raise RuntimeError("Redis client not connected")
|
||||||
|
|
||||||
|
value = await self.client.get(key)
|
||||||
|
return value.decode() if value else None
|
||||||
|
|
||||||
|
async def delete_value(self, key: str) -> None:
|
||||||
|
"""
|
||||||
|
Generic delete
|
||||||
|
"""
|
||||||
|
if not self.client:
|
||||||
|
raise RuntimeError("Redis client not connected")
|
||||||
|
|
||||||
|
await self.client.delete(key)
|
||||||
|
|
||||||
|
|
||||||
|
redis_manager = RedisManager()
|
||||||
|
|
@ -0,0 +1,261 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
SurrealDB manager with live queries for real time chat features
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
from typing import Any
|
||||||
|
from collections.abc import Callable
|
||||||
|
|
||||||
|
from surrealdb import AsyncSurreal
|
||||||
|
|
||||||
|
from app.schemas.surreal import (
|
||||||
|
LiveMessageUpdate,
|
||||||
|
LivePresenceUpdate,
|
||||||
|
MessageResponse,
|
||||||
|
PresenceResponse,
|
||||||
|
RoomResponse,
|
||||||
|
)
|
||||||
|
from app.config import DEFAULT_MESSAGE_LIMIT, settings
|
||||||
|
from app.core.enums import PresenceStatus
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class SurrealDBManager:
|
||||||
|
"""
|
||||||
|
SurrealDB connection manager with live query subscriptions
|
||||||
|
"""
|
||||||
|
def __init__(self) -> None:
|
||||||
|
"""
|
||||||
|
Initialize SurrealDB manager
|
||||||
|
"""
|
||||||
|
self.db: AsyncSurreal | None = None
|
||||||
|
self.live_queries: dict[str, str] = {}
|
||||||
|
self._connected = False
|
||||||
|
|
||||||
|
async def connect(self) -> None:
|
||||||
|
"""
|
||||||
|
Establish connection to SurrealDB
|
||||||
|
"""
|
||||||
|
if self._connected:
|
||||||
|
return
|
||||||
|
|
||||||
|
self.db = AsyncSurreal(settings.SURREAL_URL)
|
||||||
|
await self.db.connect()
|
||||||
|
|
||||||
|
await self.db.signin(
|
||||||
|
{
|
||||||
|
"username": settings.SURREAL_USER,
|
||||||
|
"password": settings.SURREAL_PASSWORD,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
await self.db.use(
|
||||||
|
settings.SURREAL_NAMESPACE,
|
||||||
|
settings.SURREAL_DATABASE,
|
||||||
|
)
|
||||||
|
|
||||||
|
self._connected = True
|
||||||
|
logger.info("Connected to SurrealDB at %s", settings.SURREAL_URL)
|
||||||
|
|
||||||
|
async def disconnect(self) -> None:
|
||||||
|
"""
|
||||||
|
Close SurrealDB connection
|
||||||
|
"""
|
||||||
|
if self.db and self._connected:
|
||||||
|
await self.db.close()
|
||||||
|
self._connected = False
|
||||||
|
logger.info("Disconnected from SurrealDB")
|
||||||
|
|
||||||
|
async def ensure_connected(self) -> None:
|
||||||
|
"""
|
||||||
|
Ensure connection is established
|
||||||
|
"""
|
||||||
|
if not self._connected:
|
||||||
|
await self.connect()
|
||||||
|
|
||||||
|
async def create_message(
|
||||||
|
self,
|
||||||
|
message_data: dict[str,
|
||||||
|
Any]
|
||||||
|
) -> MessageResponse:
|
||||||
|
"""
|
||||||
|
Create a new message in SurrealDB
|
||||||
|
"""
|
||||||
|
await self.ensure_connected()
|
||||||
|
result = await self.db.create("messages", message_data)
|
||||||
|
return MessageResponse(**result)
|
||||||
|
|
||||||
|
async def get_room_messages(
|
||||||
|
self,
|
||||||
|
room_id: str,
|
||||||
|
limit: int = DEFAULT_MESSAGE_LIMIT,
|
||||||
|
offset: int = 0,
|
||||||
|
) -> list[MessageResponse]:
|
||||||
|
"""
|
||||||
|
Get messages for a specific room with pagination
|
||||||
|
"""
|
||||||
|
await self.ensure_connected()
|
||||||
|
query = """
|
||||||
|
SELECT * FROM messages
|
||||||
|
WHERE room_id = $room_id
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
LIMIT $limit
|
||||||
|
START $offset
|
||||||
|
"""
|
||||||
|
result = await self.db.query(
|
||||||
|
query,
|
||||||
|
{
|
||||||
|
"room_id": room_id,
|
||||||
|
"limit": limit,
|
||||||
|
"offset": offset,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
messages = result[0]["result"] if result else []
|
||||||
|
return [MessageResponse(**msg) for msg in messages]
|
||||||
|
|
||||||
|
async def create_room(self, room_data: dict[str, Any]) -> RoomResponse:
|
||||||
|
"""
|
||||||
|
Create a new chat room
|
||||||
|
"""
|
||||||
|
await self.ensure_connected()
|
||||||
|
result = await self.db.create("rooms", room_data)
|
||||||
|
return RoomResponse(**result)
|
||||||
|
|
||||||
|
async def get_user_rooms(self, user_id: str) -> list[RoomResponse]:
|
||||||
|
"""
|
||||||
|
Get all rooms a user is part of using graph traversal
|
||||||
|
"""
|
||||||
|
await self.ensure_connected()
|
||||||
|
query = """
|
||||||
|
SELECT ->member_of->rooms.* AS rooms
|
||||||
|
FROM $user_id
|
||||||
|
"""
|
||||||
|
result = await self.db.query(query, {"user_id": f"users:{user_id}"})
|
||||||
|
rooms = result[0]["result"][0]["rooms"] if result else []
|
||||||
|
return [RoomResponse(**room) for room in rooms]
|
||||||
|
|
||||||
|
async def update_presence(
|
||||||
|
self,
|
||||||
|
user_id: str,
|
||||||
|
status: str,
|
||||||
|
last_seen: str,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Update user presence status
|
||||||
|
"""
|
||||||
|
await self.ensure_connected()
|
||||||
|
await self.db.merge(
|
||||||
|
f"presence:{user_id}",
|
||||||
|
{
|
||||||
|
"user_id": user_id,
|
||||||
|
"status": status,
|
||||||
|
"last_seen": last_seen,
|
||||||
|
"updated_at": "time::now()",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
async def get_room_presence(self, room_id: str) -> list[PresenceResponse]:
|
||||||
|
"""
|
||||||
|
Get presence for all users in a room
|
||||||
|
"""
|
||||||
|
await self.ensure_connected()
|
||||||
|
query = f"""
|
||||||
|
SELECT ->member_of->rooms->has_members<-presence.* AS users
|
||||||
|
FROM $room_id
|
||||||
|
WHERE status = '{PresenceStatus.ONLINE.value}'
|
||||||
|
"""
|
||||||
|
result = await self.db.query(query, {"room_id": f"rooms:{room_id}"})
|
||||||
|
presence_list = result[0]["result"] if result else []
|
||||||
|
return [PresenceResponse(**p) for p in presence_list]
|
||||||
|
|
||||||
|
async def live_messages(
|
||||||
|
self,
|
||||||
|
room_id: str,
|
||||||
|
callback: Callable[[LiveMessageUpdate],
|
||||||
|
None],
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Subscribe to live message updates for a room
|
||||||
|
"""
|
||||||
|
await self.ensure_connected()
|
||||||
|
query = f"LIVE SELECT * FROM messages WHERE room_id = '{room_id}'"
|
||||||
|
|
||||||
|
def wrapper(data: dict[str, Any]) -> None:
|
||||||
|
update = LiveMessageUpdate(**data)
|
||||||
|
callback(update)
|
||||||
|
|
||||||
|
live_id = await self.db.live(query, wrapper)
|
||||||
|
self.live_queries[room_id] = live_id
|
||||||
|
return live_id
|
||||||
|
|
||||||
|
async def live_presence(
|
||||||
|
self,
|
||||||
|
room_id: str,
|
||||||
|
callback: Callable[[LivePresenceUpdate],
|
||||||
|
None],
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Subscribe to live presence updates for a room
|
||||||
|
"""
|
||||||
|
await self.ensure_connected()
|
||||||
|
query = f"LIVE SELECT * FROM presence WHERE room_id = '{room_id}'"
|
||||||
|
|
||||||
|
def wrapper(data: dict[str, Any]) -> None:
|
||||||
|
update = LivePresenceUpdate(**data)
|
||||||
|
callback(update)
|
||||||
|
|
||||||
|
live_id = await self.db.live(query, wrapper)
|
||||||
|
self.live_queries[f"presence_{room_id}"] = live_id
|
||||||
|
return live_id
|
||||||
|
|
||||||
|
async def kill_live_query(self, live_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Stop a live query subscription
|
||||||
|
"""
|
||||||
|
await self.ensure_connected()
|
||||||
|
await self.db.kill(live_id)
|
||||||
|
|
||||||
|
for key, query_id in list(self.live_queries.items()):
|
||||||
|
if query_id == live_id:
|
||||||
|
del self.live_queries[key]
|
||||||
|
break
|
||||||
|
|
||||||
|
async def create_ephemeral_room(
|
||||||
|
self,
|
||||||
|
room_data: dict[str,
|
||||||
|
Any],
|
||||||
|
ttl_seconds: int,
|
||||||
|
) -> RoomResponse:
|
||||||
|
"""
|
||||||
|
Create an ephemeral room that auto-deletes after TTL
|
||||||
|
"""
|
||||||
|
await self.ensure_connected()
|
||||||
|
room = await self.db.create("rooms", room_data)
|
||||||
|
room_id = room["id"]
|
||||||
|
|
||||||
|
asyncio.create_task(self._schedule_room_deletion(room_id, ttl_seconds))
|
||||||
|
return RoomResponse(**room)
|
||||||
|
|
||||||
|
async def _schedule_room_deletion(
|
||||||
|
self,
|
||||||
|
room_id: str,
|
||||||
|
ttl_seconds: int
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Schedule automatic deletion of a room after TTL
|
||||||
|
"""
|
||||||
|
await asyncio.sleep(ttl_seconds)
|
||||||
|
await self.ensure_connected()
|
||||||
|
await self.db.delete(room_id)
|
||||||
|
logger.info(
|
||||||
|
"Deleted ephemeral room %s after %ss TTL",
|
||||||
|
room_id,
|
||||||
|
ttl_seconds
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
surreal_db = SurrealDBManager()
|
||||||
|
|
@ -0,0 +1,281 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
WebSocket connection manager for real time messaging
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from typing import Any
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from fastapi import WebSocket
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
WS_HEARTBEAT_INTERVAL,
|
||||||
|
WS_MAX_CONNECTIONS_PER_USER,
|
||||||
|
)
|
||||||
|
from app.core.surreal_manager import surreal_db
|
||||||
|
from app.schemas.surreal import LiveMessageUpdate
|
||||||
|
from app.schemas.websocket import (
|
||||||
|
EncryptedMessageWS,
|
||||||
|
ErrorMessageWS,
|
||||||
|
WSHeartbeat,
|
||||||
|
)
|
||||||
|
from app.services.presence_service import presence_service
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class ConnectionManager:
|
||||||
|
"""
|
||||||
|
Manages WebSocket connections and message broadcasting
|
||||||
|
"""
|
||||||
|
def __init__(self) -> None:
|
||||||
|
"""
|
||||||
|
Initialize connection manager with empty connection pool
|
||||||
|
"""
|
||||||
|
self.active_connections: dict[UUID, list[WebSocket]] = {}
|
||||||
|
self.live_query_ids: dict[UUID, str] = {}
|
||||||
|
self.heartbeat_tasks: dict[UUID, asyncio.Task] = {}
|
||||||
|
|
||||||
|
async def connect(self, websocket: WebSocket, user_id: UUID) -> bool:
|
||||||
|
"""
|
||||||
|
Accept WebSocket connection and register user
|
||||||
|
"""
|
||||||
|
await websocket.accept()
|
||||||
|
|
||||||
|
if user_id not in self.active_connections:
|
||||||
|
self.active_connections[user_id] = []
|
||||||
|
|
||||||
|
if len(self.active_connections[user_id]) >= WS_MAX_CONNECTIONS_PER_USER:
|
||||||
|
logger.warning(
|
||||||
|
"User %s exceeded max connections (%s)",
|
||||||
|
user_id,
|
||||||
|
WS_MAX_CONNECTIONS_PER_USER
|
||||||
|
)
|
||||||
|
await self._send_error(
|
||||||
|
websocket,
|
||||||
|
"max_connections",
|
||||||
|
f"Maximum {WS_MAX_CONNECTIONS_PER_USER} connections per user"
|
||||||
|
)
|
||||||
|
await websocket.close()
|
||||||
|
return False
|
||||||
|
|
||||||
|
self.active_connections[user_id].append(websocket)
|
||||||
|
logger.info(
|
||||||
|
"User %s connected via WebSocket (total: %s)",
|
||||||
|
user_id,
|
||||||
|
len(self.active_connections[user_id])
|
||||||
|
)
|
||||||
|
|
||||||
|
await presence_service.set_user_online(user_id)
|
||||||
|
|
||||||
|
self.heartbeat_tasks[user_id] = asyncio.create_task(
|
||||||
|
self._heartbeat_loop(websocket,
|
||||||
|
user_id)
|
||||||
|
)
|
||||||
|
|
||||||
|
await self._subscribe_to_messages(user_id)
|
||||||
|
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def disconnect(self, websocket: WebSocket, user_id: UUID) -> None:
|
||||||
|
"""
|
||||||
|
Remove WebSocket connection and cleanup resources
|
||||||
|
"""
|
||||||
|
if user_id in self.active_connections:
|
||||||
|
if websocket in self.active_connections[user_id]:
|
||||||
|
self.active_connections[user_id].remove(websocket)
|
||||||
|
|
||||||
|
if not self.active_connections[user_id]:
|
||||||
|
del self.active_connections[user_id]
|
||||||
|
|
||||||
|
await presence_service.set_user_offline(user_id)
|
||||||
|
|
||||||
|
if user_id in self.live_query_ids:
|
||||||
|
try:
|
||||||
|
await surreal_db.kill_live_query(
|
||||||
|
self.live_query_ids[user_id]
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
"Failed to kill live query for %s: %s",
|
||||||
|
user_id,
|
||||||
|
e
|
||||||
|
)
|
||||||
|
del self.live_query_ids[user_id]
|
||||||
|
|
||||||
|
if user_id in self.heartbeat_tasks:
|
||||||
|
self.heartbeat_tasks[user_id].cancel()
|
||||||
|
del self.heartbeat_tasks[user_id]
|
||||||
|
|
||||||
|
logger.info("User %s fully disconnected", user_id)
|
||||||
|
else:
|
||||||
|
logger.info(
|
||||||
|
"User %s connection closed (remaining: %s)",
|
||||||
|
user_id,
|
||||||
|
len(self.active_connections[user_id])
|
||||||
|
)
|
||||||
|
|
||||||
|
async def send_message(self, user_id: UUID, message: dict[str, Any]) -> None:
|
||||||
|
"""
|
||||||
|
Send message to all connections for a specific user
|
||||||
|
"""
|
||||||
|
if user_id not in self.active_connections:
|
||||||
|
logger.debug("User %s not connected, cannot send message", user_id)
|
||||||
|
return
|
||||||
|
|
||||||
|
dead_connections = []
|
||||||
|
|
||||||
|
for websocket in self.active_connections[user_id]:
|
||||||
|
try:
|
||||||
|
await websocket.send_json(message)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to send message to %s: %s", user_id, e)
|
||||||
|
dead_connections.append(websocket)
|
||||||
|
|
||||||
|
for dead_ws in dead_connections:
|
||||||
|
await self.disconnect(dead_ws, user_id)
|
||||||
|
|
||||||
|
async def broadcast_to_room(
|
||||||
|
self,
|
||||||
|
room_id: str,
|
||||||
|
message: dict[str,
|
||||||
|
Any]
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Broadcast message to all users in a room
|
||||||
|
"""
|
||||||
|
online_users = await presence_service.get_room_online_users(room_id)
|
||||||
|
|
||||||
|
for user_presence in online_users:
|
||||||
|
try:
|
||||||
|
user_id = UUID(user_presence["user_id"])
|
||||||
|
await self.send_message(user_id, message)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
"Failed to broadcast to user %s: %s",
|
||||||
|
user_presence['user_id'],
|
||||||
|
e
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _heartbeat_loop(self, websocket: WebSocket, user_id: UUID) -> None:
|
||||||
|
"""
|
||||||
|
Send periodic heartbeat pings to keep connection alive
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
await asyncio.sleep(WS_HEARTBEAT_INTERVAL)
|
||||||
|
|
||||||
|
if user_id not in self.active_connections:
|
||||||
|
break
|
||||||
|
|
||||||
|
if websocket not in self.active_connections[user_id]:
|
||||||
|
break
|
||||||
|
|
||||||
|
heartbeat = WSHeartbeat(timestamp = datetime.now(UTC))
|
||||||
|
|
||||||
|
try:
|
||||||
|
await websocket.send_json(heartbeat.model_dump(mode = "json"))
|
||||||
|
await presence_service.update_last_seen(user_id)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Heartbeat failed for user %s: %s", user_id, e)
|
||||||
|
await self.disconnect(websocket, user_id)
|
||||||
|
break
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
logger.debug("Heartbeat task cancelled for user %s", user_id)
|
||||||
|
|
||||||
|
async def _subscribe_to_messages(self, user_id: UUID) -> None:
|
||||||
|
"""
|
||||||
|
Subscribe to live message updates for the user
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
|
||||||
|
def message_callback(update: LiveMessageUpdate) -> None:
|
||||||
|
"""
|
||||||
|
Handle incoming message from SurrealDB live query
|
||||||
|
"""
|
||||||
|
asyncio.create_task(self._handle_live_message(user_id, update))
|
||||||
|
|
||||||
|
live_id = await surreal_db.live_messages(
|
||||||
|
room_id = str(user_id),
|
||||||
|
callback = message_callback
|
||||||
|
)
|
||||||
|
|
||||||
|
self.live_query_ids[user_id] = live_id
|
||||||
|
logger.debug("Subscribed to live messages for user %s", user_id)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to subscribe to messages for %s: %s", user_id, e)
|
||||||
|
|
||||||
|
async def _handle_live_message(
|
||||||
|
self,
|
||||||
|
user_id: UUID,
|
||||||
|
update: LiveMessageUpdate
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Process live message update and forward to WebSocket
|
||||||
|
"""
|
||||||
|
if update.action != "CREATE":
|
||||||
|
return
|
||||||
|
|
||||||
|
message_data = update.result
|
||||||
|
|
||||||
|
ws_message = EncryptedMessageWS(
|
||||||
|
message_id = message_data.id,
|
||||||
|
sender_id = message_data.sender_id,
|
||||||
|
recipient_id = str(user_id),
|
||||||
|
ciphertext = message_data.encrypted_content,
|
||||||
|
nonce = "",
|
||||||
|
header = message_data.encrypted_header,
|
||||||
|
sender_username = "",
|
||||||
|
timestamp = message_data.created_at
|
||||||
|
)
|
||||||
|
|
||||||
|
await self.send_message(user_id, ws_message.model_dump(mode = "json"))
|
||||||
|
|
||||||
|
async def _send_error(
|
||||||
|
self,
|
||||||
|
websocket: WebSocket,
|
||||||
|
error_code: str,
|
||||||
|
error_message: str
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Send error message to WebSocket connection
|
||||||
|
"""
|
||||||
|
error = ErrorMessageWS(
|
||||||
|
error_code = error_code,
|
||||||
|
error_message = error_message,
|
||||||
|
timestamp = datetime.now(UTC)
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await websocket.send_json(error.model_dump(mode = "json"))
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to send error message: %s", e)
|
||||||
|
|
||||||
|
def get_active_users(self) -> list[UUID]:
|
||||||
|
"""
|
||||||
|
Get list of all currently connected user IDs
|
||||||
|
"""
|
||||||
|
return list(self.active_connections.keys())
|
||||||
|
|
||||||
|
def get_connection_count(self, user_id: UUID) -> int:
|
||||||
|
"""
|
||||||
|
Get number of active connections for a user
|
||||||
|
"""
|
||||||
|
if user_id not in self.active_connections:
|
||||||
|
return 0
|
||||||
|
return len(self.active_connections[user_id])
|
||||||
|
|
||||||
|
def is_user_connected(self, user_id: UUID) -> bool:
|
||||||
|
"""
|
||||||
|
Check if user has any active connections
|
||||||
|
"""
|
||||||
|
return user_id in self.active_connections and len(
|
||||||
|
self.active_connections[user_id]
|
||||||
|
) > 0
|
||||||
|
|
||||||
|
|
||||||
|
connection_manager = ConnectionManager()
|
||||||
|
|
@ -0,0 +1,112 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
FastAPI application factory
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from collections.abc import AsyncGenerator
|
||||||
|
from contextlib import asynccontextmanager
|
||||||
|
|
||||||
|
from fastapi import FastAPI
|
||||||
|
from fastapi.responses import ORJSONResponse
|
||||||
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
|
from fastapi.middleware.gzip import GZipMiddleware
|
||||||
|
|
||||||
|
from app.schemas.common import (
|
||||||
|
HealthResponse,
|
||||||
|
RootResponse,
|
||||||
|
)
|
||||||
|
from app.config import (
|
||||||
|
settings,
|
||||||
|
APP_VERSION,
|
||||||
|
APP_STATUS,
|
||||||
|
APP_DESCRIPTION,
|
||||||
|
GZIP_MINIMUM_SIZE,
|
||||||
|
)
|
||||||
|
from app.models.Base import init_db
|
||||||
|
from app.api.auth import router as auth_router
|
||||||
|
from app.core.surreal_manager import surreal_db
|
||||||
|
from app.core.redis_manager import redis_manager
|
||||||
|
from app.api.encryption import router as encryption_router
|
||||||
|
from app.api.websocket import router as websocket_router
|
||||||
|
from app.core.exception_handlers import register_exception_handlers
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def lifespan(app: FastAPI) -> AsyncGenerator[None]:
|
||||||
|
"""
|
||||||
|
Application lifespan manager for startup and shutdown events
|
||||||
|
"""
|
||||||
|
logger.info("Starting %s in %s mode", settings.APP_NAME, settings.ENV)
|
||||||
|
|
||||||
|
await init_db()
|
||||||
|
logger.info("PostgreSQL database initialized")
|
||||||
|
|
||||||
|
await redis_manager.connect()
|
||||||
|
logger.info("Redis connected")
|
||||||
|
|
||||||
|
await surreal_db.connect()
|
||||||
|
logger.info("SurrealDB connected")
|
||||||
|
|
||||||
|
yield
|
||||||
|
|
||||||
|
logger.info("Shutting down application")
|
||||||
|
|
||||||
|
await redis_manager.disconnect()
|
||||||
|
await surreal_db.disconnect()
|
||||||
|
|
||||||
|
|
||||||
|
def create_app() -> FastAPI:
|
||||||
|
"""
|
||||||
|
Create and configure the FastAPI application instance
|
||||||
|
"""
|
||||||
|
app = FastAPI(
|
||||||
|
title = settings.APP_NAME,
|
||||||
|
description = APP_DESCRIPTION,
|
||||||
|
version = APP_VERSION,
|
||||||
|
docs_url = "/docs" if settings.is_development else None,
|
||||||
|
redoc_url = "/redoc" if settings.is_development else None,
|
||||||
|
default_response_class = ORJSONResponse,
|
||||||
|
lifespan = lifespan,
|
||||||
|
)
|
||||||
|
|
||||||
|
app.add_middleware(
|
||||||
|
CORSMiddleware,
|
||||||
|
allow_origins = settings.CORS_ORIGINS,
|
||||||
|
allow_credentials = True,
|
||||||
|
allow_methods = ["*"],
|
||||||
|
allow_headers = ["*"],
|
||||||
|
expose_headers = ["*"],
|
||||||
|
)
|
||||||
|
|
||||||
|
app.add_middleware(GZipMiddleware, minimum_size = GZIP_MINIMUM_SIZE)
|
||||||
|
|
||||||
|
register_exception_handlers(app)
|
||||||
|
|
||||||
|
@app.get("/", tags = ["root"])
|
||||||
|
async def root() -> RootResponse:
|
||||||
|
"""
|
||||||
|
Root endpoint returning API status
|
||||||
|
"""
|
||||||
|
return RootResponse(
|
||||||
|
app = settings.APP_NAME,
|
||||||
|
version = APP_VERSION,
|
||||||
|
status = APP_STATUS,
|
||||||
|
environment = settings.ENV,
|
||||||
|
)
|
||||||
|
|
||||||
|
@app.get("/health", tags = ["health"])
|
||||||
|
async def health() -> HealthResponse:
|
||||||
|
"""
|
||||||
|
Health check endpoint for monitoring
|
||||||
|
"""
|
||||||
|
return HealthResponse(status = "healthy")
|
||||||
|
|
||||||
|
app.include_router(auth_router)
|
||||||
|
app.include_router(encryption_router)
|
||||||
|
app.include_router(websocket_router)
|
||||||
|
|
||||||
|
return app
|
||||||
|
|
@ -0,0 +1,31 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Application entry point with uvicorn server command
|
||||||
|
"""
|
||||||
|
|
||||||
|
import uvicorn
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
DEFAULT_HOST,
|
||||||
|
DEFAULT_PORT,
|
||||||
|
settings,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
"""
|
||||||
|
Run the FastAPI application with uvicorn
|
||||||
|
"""
|
||||||
|
uvicorn.run(
|
||||||
|
"app.factory:create_app",
|
||||||
|
factory = True,
|
||||||
|
host = DEFAULT_HOST,
|
||||||
|
port = DEFAULT_PORT,
|
||||||
|
reload = settings.is_development,
|
||||||
|
log_level = "debug" if settings.DEBUG else "info",
|
||||||
|
access_log = True,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
|
@ -0,0 +1,67 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Base SQLModel class with async PostgreSQL engine setup
|
||||||
|
"""
|
||||||
|
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from collections.abc import AsyncGenerator
|
||||||
|
|
||||||
|
from sqlalchemy import DateTime
|
||||||
|
from sqlalchemy.ext.asyncio import (
|
||||||
|
AsyncSession,
|
||||||
|
create_async_engine,
|
||||||
|
)
|
||||||
|
from sqlmodel import Field, SQLModel
|
||||||
|
from sqlalchemy.orm import sessionmaker
|
||||||
|
|
||||||
|
from app.config import settings
|
||||||
|
|
||||||
|
|
||||||
|
class BaseDBModel(SQLModel):
|
||||||
|
"""
|
||||||
|
Base model with common timestamp fields
|
||||||
|
"""
|
||||||
|
created_at: datetime = Field(
|
||||||
|
default_factory = lambda: datetime.now(UTC),
|
||||||
|
nullable = False,
|
||||||
|
sa_type = DateTime(timezone = True),
|
||||||
|
)
|
||||||
|
updated_at: datetime = Field(
|
||||||
|
default_factory = lambda: datetime.now(UTC),
|
||||||
|
nullable = False,
|
||||||
|
sa_type = DateTime(timezone = True),
|
||||||
|
sa_column_kwargs = {"onupdate": lambda: datetime.now(UTC)},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# Create async engine for PostgreSQL
|
||||||
|
engine = create_async_engine(
|
||||||
|
str(settings.DATABASE_URL),
|
||||||
|
echo = settings.DEBUG,
|
||||||
|
pool_size = settings.DB_POOL_SIZE,
|
||||||
|
max_overflow = settings.DB_MAX_OVERFLOW,
|
||||||
|
pool_pre_ping = True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create async session factory
|
||||||
|
async_session_maker = sessionmaker( # type: ignore[call-overload]
|
||||||
|
bind = engine,
|
||||||
|
class_ = AsyncSession,
|
||||||
|
expire_on_commit = False,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def get_session() -> AsyncGenerator[AsyncSession]:
|
||||||
|
"""
|
||||||
|
Dependency for getting database sessions
|
||||||
|
"""
|
||||||
|
async with async_session_maker() as session:
|
||||||
|
yield session
|
||||||
|
|
||||||
|
|
||||||
|
async def init_db() -> None:
|
||||||
|
"""
|
||||||
|
Initialize database tables
|
||||||
|
"""
|
||||||
|
async with engine.begin() as conn:
|
||||||
|
await conn.run_sync(SQLModel.metadata.create_all)
|
||||||
|
|
@ -0,0 +1,78 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
WebAuthn credential model for passkey storage
|
||||||
|
"""
|
||||||
|
|
||||||
|
from datetime import datetime
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from sqlalchemy import DateTime
|
||||||
|
from sqlmodel import Field, Relationship
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
AAGUID_MAX_LENGTH,
|
||||||
|
ATTESTATION_TYPE_MAX_LENGTH,
|
||||||
|
CREDENTIAL_ID_MAX_LENGTH,
|
||||||
|
DISPLAY_NAME_MAX_LENGTH,
|
||||||
|
PUBLIC_KEY_MAX_LENGTH,
|
||||||
|
TRANSPORT_MAX_LENGTH,
|
||||||
|
)
|
||||||
|
from app.models.Base import BaseDBModel
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from app.models.User import User
|
||||||
|
|
||||||
|
|
||||||
|
class Credential(BaseDBModel, table = True):
|
||||||
|
"""
|
||||||
|
WebAuthn/FIDO2 passkey credential
|
||||||
|
"""
|
||||||
|
__tablename__ = "credentials"
|
||||||
|
|
||||||
|
id: int = Field(default = None, primary_key = True)
|
||||||
|
credential_id: str = Field(
|
||||||
|
unique = True,
|
||||||
|
index = True,
|
||||||
|
nullable = False,
|
||||||
|
max_length = CREDENTIAL_ID_MAX_LENGTH
|
||||||
|
)
|
||||||
|
public_key: str = Field(nullable = False, max_length = PUBLIC_KEY_MAX_LENGTH)
|
||||||
|
sign_count: int = Field(default = 0, nullable = False)
|
||||||
|
aaguid: str | None = Field(default = None, max_length = AAGUID_MAX_LENGTH)
|
||||||
|
|
||||||
|
# WebAuthn Level 3 fields
|
||||||
|
backup_eligible: bool = Field(default = False, nullable = False)
|
||||||
|
backup_state: bool = Field(default = False, nullable = False)
|
||||||
|
attestation_type: str | None = Field(
|
||||||
|
default = None,
|
||||||
|
max_length = ATTESTATION_TYPE_MAX_LENGTH
|
||||||
|
)
|
||||||
|
transports: str | None = Field(
|
||||||
|
default = None,
|
||||||
|
max_length = TRANSPORT_MAX_LENGTH
|
||||||
|
)
|
||||||
|
|
||||||
|
# User relationship
|
||||||
|
user_id: UUID = Field(
|
||||||
|
foreign_key = "users.id",
|
||||||
|
nullable = False,
|
||||||
|
index = True
|
||||||
|
)
|
||||||
|
user: "User" = Relationship(back_populates = "credentials")
|
||||||
|
|
||||||
|
# Device metadata
|
||||||
|
device_name: str | None = Field(
|
||||||
|
default = None,
|
||||||
|
max_length = DISPLAY_NAME_MAX_LENGTH
|
||||||
|
)
|
||||||
|
last_used_at: datetime | None = Field(
|
||||||
|
default = None,
|
||||||
|
sa_type = DateTime(timezone = True),
|
||||||
|
)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
"""
|
||||||
|
String representation of Credential
|
||||||
|
"""
|
||||||
|
return f"<Credential {self.credential_id[:16]}...>"
|
||||||
|
|
@ -0,0 +1,48 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
X3DH identity key model for long term user identification
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from sqlmodel import Field
|
||||||
|
|
||||||
|
from app.config import IDENTITY_KEY_LENGTH
|
||||||
|
from app.models.Base import BaseDBModel
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class IdentityKey(BaseDBModel, table = True):
|
||||||
|
"""
|
||||||
|
Long term X25519 identity key for X3DH protocol
|
||||||
|
"""
|
||||||
|
__tablename__ = "identity_keys"
|
||||||
|
|
||||||
|
id: int = Field(default = None, primary_key = True)
|
||||||
|
user_id: UUID = Field(
|
||||||
|
foreign_key = "users.id",
|
||||||
|
nullable = False,
|
||||||
|
unique = True,
|
||||||
|
index = True
|
||||||
|
)
|
||||||
|
|
||||||
|
public_key: str = Field(nullable = False, max_length = IDENTITY_KEY_LENGTH)
|
||||||
|
private_key: str = Field(nullable = False, max_length = IDENTITY_KEY_LENGTH)
|
||||||
|
|
||||||
|
public_key_ed25519: str = Field(
|
||||||
|
nullable = False,
|
||||||
|
max_length = IDENTITY_KEY_LENGTH
|
||||||
|
)
|
||||||
|
private_key_ed25519: str = Field(
|
||||||
|
nullable = False,
|
||||||
|
max_length = IDENTITY_KEY_LENGTH
|
||||||
|
)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
"""
|
||||||
|
String representation of IdentityKey
|
||||||
|
"""
|
||||||
|
return f"<IdentityKey user_id={self.user_id}>"
|
||||||
|
|
@ -0,0 +1,45 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
X3DH one time prekey model for single use key exchange
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from sqlmodel import Field
|
||||||
|
|
||||||
|
from app.config import ONE_TIME_PREKEY_LENGTH
|
||||||
|
from app.models.Base import BaseDBModel
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class OneTimePrekey(BaseDBModel, table = True):
|
||||||
|
"""
|
||||||
|
X25519 one time prekey consumed after single use for X3DH protocol
|
||||||
|
"""
|
||||||
|
__tablename__ = "one_time_prekeys"
|
||||||
|
|
||||||
|
id: int = Field(default = None, primary_key = True)
|
||||||
|
user_id: UUID = Field(
|
||||||
|
foreign_key = "users.id",
|
||||||
|
nullable = False,
|
||||||
|
index = True
|
||||||
|
)
|
||||||
|
|
||||||
|
key_id: int = Field(nullable = False, index = True)
|
||||||
|
|
||||||
|
public_key: str = Field(nullable = False, max_length = ONE_TIME_PREKEY_LENGTH)
|
||||||
|
private_key: str = Field(
|
||||||
|
nullable = False,
|
||||||
|
max_length = ONE_TIME_PREKEY_LENGTH
|
||||||
|
)
|
||||||
|
|
||||||
|
is_used: bool = Field(default = False, nullable = False, index = True)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
"""
|
||||||
|
String representation of OneTimePrekey
|
||||||
|
"""
|
||||||
|
return f"<OneTimePrekey user_id={self.user_id} key_id={self.key_id}>"
|
||||||
|
|
@ -0,0 +1,68 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Double Ratchet state model for per conversation encryption state
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from sqlmodel import Field
|
||||||
|
|
||||||
|
from app.config import RATCHET_STATE_MAX_LENGTH
|
||||||
|
from app.models.Base import BaseDBModel
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class RatchetState(BaseDBModel, table = True):
|
||||||
|
"""
|
||||||
|
Double Ratchet algorithm state for a conversation between two users
|
||||||
|
"""
|
||||||
|
__tablename__ = "ratchet_states"
|
||||||
|
|
||||||
|
id: int = Field(default = None, primary_key = True)
|
||||||
|
|
||||||
|
user_id: UUID = Field(
|
||||||
|
foreign_key = "users.id",
|
||||||
|
nullable = False,
|
||||||
|
index = True
|
||||||
|
)
|
||||||
|
peer_user_id: UUID = Field(
|
||||||
|
foreign_key = "users.id",
|
||||||
|
nullable = False,
|
||||||
|
index = True
|
||||||
|
)
|
||||||
|
|
||||||
|
dh_private_key: str | None = Field(
|
||||||
|
default = None,
|
||||||
|
max_length = RATCHET_STATE_MAX_LENGTH
|
||||||
|
)
|
||||||
|
dh_public_key: str | None = Field(
|
||||||
|
default = None,
|
||||||
|
max_length = RATCHET_STATE_MAX_LENGTH
|
||||||
|
)
|
||||||
|
dh_peer_public_key: str | None = Field(
|
||||||
|
default = None,
|
||||||
|
max_length = RATCHET_STATE_MAX_LENGTH
|
||||||
|
)
|
||||||
|
|
||||||
|
root_key: str = Field(nullable = False, max_length = RATCHET_STATE_MAX_LENGTH)
|
||||||
|
sending_chain_key: str = Field(
|
||||||
|
nullable = False,
|
||||||
|
max_length = RATCHET_STATE_MAX_LENGTH
|
||||||
|
)
|
||||||
|
receiving_chain_key: str = Field(
|
||||||
|
nullable = False,
|
||||||
|
max_length = RATCHET_STATE_MAX_LENGTH
|
||||||
|
)
|
||||||
|
|
||||||
|
sending_message_number: int = Field(default = 0, nullable = False)
|
||||||
|
receiving_message_number: int = Field(default = 0, nullable = False)
|
||||||
|
previous_sending_chain_length: int = Field(default = 0, nullable = False)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
"""
|
||||||
|
String representation of RatchetState
|
||||||
|
"""
|
||||||
|
return f"<RatchetState user_id={self.user_id} peer={self.peer_user_id}>"
|
||||||
|
|
@ -0,0 +1,50 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
X3DH signed prekey model for medium term key rotation
|
||||||
|
"""
|
||||||
|
|
||||||
|
from datetime import datetime
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from sqlalchemy import DateTime
|
||||||
|
from sqlmodel import Field
|
||||||
|
|
||||||
|
from app.config import SIGNATURE_LENGTH, SIGNED_PREKEY_LENGTH
|
||||||
|
from app.models.Base import BaseDBModel
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class SignedPrekey(BaseDBModel, table = True):
|
||||||
|
"""
|
||||||
|
X25519 signed prekey rotated every 48 hours for X3DH protocol
|
||||||
|
"""
|
||||||
|
__tablename__ = "signed_prekeys"
|
||||||
|
|
||||||
|
id: int = Field(default = None, primary_key = True)
|
||||||
|
user_id: UUID = Field(
|
||||||
|
foreign_key = "users.id",
|
||||||
|
nullable = False,
|
||||||
|
index = True
|
||||||
|
)
|
||||||
|
|
||||||
|
key_id: int = Field(nullable = False, index = True)
|
||||||
|
|
||||||
|
public_key: str = Field(nullable = False, max_length = SIGNED_PREKEY_LENGTH)
|
||||||
|
private_key: str = Field(nullable = False, max_length = SIGNED_PREKEY_LENGTH)
|
||||||
|
|
||||||
|
signature: str = Field(nullable = False, max_length = SIGNATURE_LENGTH)
|
||||||
|
|
||||||
|
is_active: bool = Field(default = True, nullable = False)
|
||||||
|
expires_at: datetime | None = Field(
|
||||||
|
default = None,
|
||||||
|
sa_type = DateTime(timezone = True),
|
||||||
|
)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
"""
|
||||||
|
String representation of SignedPrekey
|
||||||
|
"""
|
||||||
|
return f"<SignedPrekey user_id={self.user_id} key_id={self.key_id}>"
|
||||||
|
|
@ -0,0 +1,51 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Skipped message key storage for out of order Double Ratchet messages
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from sqlmodel import Field
|
||||||
|
|
||||||
|
from app.models.Base import BaseDBModel
|
||||||
|
from app.config import RATCHET_STATE_MAX_LENGTH
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class SkippedMessageKey(BaseDBModel, table = True):
|
||||||
|
"""
|
||||||
|
Stores message keys for out of order messages in Double Ratchet
|
||||||
|
"""
|
||||||
|
__tablename__ = "skipped_message_keys"
|
||||||
|
|
||||||
|
id: int = Field(default = None, primary_key = True)
|
||||||
|
|
||||||
|
ratchet_state_id: int = Field(
|
||||||
|
foreign_key = "ratchet_states.id",
|
||||||
|
nullable = False,
|
||||||
|
index = True
|
||||||
|
)
|
||||||
|
|
||||||
|
dh_public_key: str = Field(
|
||||||
|
nullable = False,
|
||||||
|
max_length = RATCHET_STATE_MAX_LENGTH,
|
||||||
|
index = True
|
||||||
|
)
|
||||||
|
message_number: int = Field(nullable = False, index = True)
|
||||||
|
|
||||||
|
message_key: str = Field(
|
||||||
|
nullable = False,
|
||||||
|
max_length = RATCHET_STATE_MAX_LENGTH
|
||||||
|
)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
"""
|
||||||
|
String representation of SkippedMessageKey
|
||||||
|
"""
|
||||||
|
return (
|
||||||
|
f"<SkippedMessageKey "
|
||||||
|
f"ratchet_id={self.ratchet_state_id} "
|
||||||
|
f"msg_num={self.message_number}>"
|
||||||
|
)
|
||||||
|
|
@ -0,0 +1,68 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
User model for authentication stored in PostgreSQL
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
from uuid import UUID, uuid4
|
||||||
|
|
||||||
|
from sqlmodel import (
|
||||||
|
Field,
|
||||||
|
Relationship,
|
||||||
|
)
|
||||||
|
from app.config import (
|
||||||
|
DISPLAY_NAME_MAX_LENGTH,
|
||||||
|
PREKEY_MAX_LENGTH,
|
||||||
|
USERNAME_MAX_LENGTH,
|
||||||
|
)
|
||||||
|
from app.models.Base import BaseDBModel
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from app.models.Credential import Credential
|
||||||
|
|
||||||
|
|
||||||
|
class User(BaseDBModel, table = True):
|
||||||
|
"""
|
||||||
|
User account with WebAuthn passkey authentication
|
||||||
|
"""
|
||||||
|
__tablename__ = "users"
|
||||||
|
|
||||||
|
id: UUID = Field(
|
||||||
|
default_factory = uuid4,
|
||||||
|
primary_key = True,
|
||||||
|
nullable = False
|
||||||
|
)
|
||||||
|
username: str = Field(
|
||||||
|
unique = True,
|
||||||
|
index = True,
|
||||||
|
nullable = False,
|
||||||
|
max_length = USERNAME_MAX_LENGTH
|
||||||
|
)
|
||||||
|
display_name: str = Field(
|
||||||
|
nullable = False,
|
||||||
|
max_length = DISPLAY_NAME_MAX_LENGTH
|
||||||
|
)
|
||||||
|
is_active: bool = Field(default = True, nullable = False)
|
||||||
|
is_verified: bool = Field(default = False, nullable = False)
|
||||||
|
|
||||||
|
credentials: list["Credential"] = Relationship(back_populates = "user")
|
||||||
|
|
||||||
|
identity_key: str | None = Field(
|
||||||
|
default = None,
|
||||||
|
max_length = PREKEY_MAX_LENGTH
|
||||||
|
)
|
||||||
|
signed_prekey: str | None = Field(
|
||||||
|
default = None,
|
||||||
|
max_length = PREKEY_MAX_LENGTH
|
||||||
|
)
|
||||||
|
signed_prekey_signature: str | None = Field(
|
||||||
|
default = None,
|
||||||
|
max_length = PREKEY_MAX_LENGTH
|
||||||
|
)
|
||||||
|
one_time_prekeys: str | None = Field(default = None)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
"""
|
||||||
|
String representation of User
|
||||||
|
"""
|
||||||
|
return f"<User {self.username}>"
|
||||||
|
|
@ -0,0 +1,33 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Database models exports
|
||||||
|
"""
|
||||||
|
|
||||||
|
from app.models.Base import (
|
||||||
|
BaseDBModel,
|
||||||
|
engine,
|
||||||
|
get_session,
|
||||||
|
init_db,
|
||||||
|
)
|
||||||
|
from app.models.Credential import Credential
|
||||||
|
from app.models.IdentityKey import IdentityKey
|
||||||
|
from app.models.OneTimePrekey import OneTimePrekey
|
||||||
|
from app.models.RatchetState import RatchetState
|
||||||
|
from app.models.SignedPrekey import SignedPrekey
|
||||||
|
from app.models.SkippedMessageKey import SkippedMessageKey
|
||||||
|
from app.models.User import User
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"BaseDBModel",
|
||||||
|
"Credential",
|
||||||
|
"IdentityKey",
|
||||||
|
"OneTimePrekey",
|
||||||
|
"RatchetState",
|
||||||
|
"SignedPrekey",
|
||||||
|
"SkippedMessageKey",
|
||||||
|
"User",
|
||||||
|
"engine",
|
||||||
|
"get_session",
|
||||||
|
"init_db",
|
||||||
|
]
|
||||||
|
|
@ -0,0 +1,64 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Pydantic schemas exports
|
||||||
|
"""
|
||||||
|
|
||||||
|
from app.schemas.auth import (
|
||||||
|
AuthenticationBeginRequest,
|
||||||
|
AuthenticationCompleteRequest,
|
||||||
|
AuthenticationOptionsResponse,
|
||||||
|
RegistrationBeginRequest,
|
||||||
|
RegistrationCompleteRequest,
|
||||||
|
RegistrationOptionsResponse,
|
||||||
|
UserResponse,
|
||||||
|
VerifiedAuthentication,
|
||||||
|
VerifiedRegistration,
|
||||||
|
)
|
||||||
|
from app.schemas.surreal import (
|
||||||
|
LiveMessageUpdate,
|
||||||
|
LivePresenceUpdate,
|
||||||
|
LiveQueryUpdate,
|
||||||
|
MessageResponse,
|
||||||
|
PresenceResponse,
|
||||||
|
RoomResponse,
|
||||||
|
)
|
||||||
|
from app.schemas.websocket import (
|
||||||
|
BaseWSMessage,
|
||||||
|
EncryptedMessageWS,
|
||||||
|
ErrorMessageWS,
|
||||||
|
PresenceUpdateWS,
|
||||||
|
ReadReceiptWS,
|
||||||
|
TypingIndicatorWS,
|
||||||
|
WSConnectionRequest,
|
||||||
|
WSHeartbeat,
|
||||||
|
)
|
||||||
|
from app.schemas.common import HealthResponse, RootResponse
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"MessageResponse",
|
||||||
|
"RoomResponse",
|
||||||
|
"PresenceResponse",
|
||||||
|
"LiveQueryUpdate",
|
||||||
|
"LiveMessageUpdate",
|
||||||
|
"LivePresenceUpdate",
|
||||||
|
"RegistrationOptionsResponse",
|
||||||
|
"VerifiedRegistration",
|
||||||
|
"AuthenticationOptionsResponse",
|
||||||
|
"VerifiedAuthentication",
|
||||||
|
"RegistrationBeginRequest",
|
||||||
|
"RegistrationCompleteRequest",
|
||||||
|
"AuthenticationBeginRequest",
|
||||||
|
"AuthenticationCompleteRequest",
|
||||||
|
"UserResponse",
|
||||||
|
"BaseWSMessage",
|
||||||
|
"EncryptedMessageWS",
|
||||||
|
"TypingIndicatorWS",
|
||||||
|
"PresenceUpdateWS",
|
||||||
|
"ReadReceiptWS",
|
||||||
|
"ErrorMessageWS",
|
||||||
|
"WSConnectionRequest",
|
||||||
|
"WSHeartbeat",
|
||||||
|
"RootResponse",
|
||||||
|
"HealthResponse",
|
||||||
|
]
|
||||||
|
|
@ -0,0 +1,142 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Pydantic schemas for WebAuthn authentication
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
DEVICE_NAME_MAX_LENGTH,
|
||||||
|
DISPLAY_NAME_MAX_LENGTH,
|
||||||
|
DISPLAY_NAME_MIN_LENGTH,
|
||||||
|
USERNAME_MAX_LENGTH,
|
||||||
|
USERNAME_MIN_LENGTH,
|
||||||
|
USER_SEARCH_DEFAULT_LIMIT,
|
||||||
|
USER_SEARCH_MAX_LIMIT,
|
||||||
|
USER_SEARCH_MIN_LENGTH,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class RegistrationOptionsResponse(BaseModel):
|
||||||
|
"""
|
||||||
|
WebAuthn registration options returned to client
|
||||||
|
"""
|
||||||
|
options: dict[str, Any]
|
||||||
|
challenge: bytes
|
||||||
|
|
||||||
|
|
||||||
|
class VerifiedRegistration(BaseModel):
|
||||||
|
"""
|
||||||
|
Verified WebAuthn registration data
|
||||||
|
"""
|
||||||
|
credential_id: bytes
|
||||||
|
credential_public_key: bytes
|
||||||
|
sign_count: int
|
||||||
|
aaguid: bytes
|
||||||
|
attestation_object: bytes
|
||||||
|
credential_type: str
|
||||||
|
user_verified: bool
|
||||||
|
attestation_format: str
|
||||||
|
credential_device_type: str
|
||||||
|
credential_backed_up: bool
|
||||||
|
backup_eligible: bool
|
||||||
|
backup_state: bool
|
||||||
|
|
||||||
|
|
||||||
|
class AuthenticationOptionsResponse(BaseModel):
|
||||||
|
"""
|
||||||
|
WebAuthn authentication options returned to client
|
||||||
|
"""
|
||||||
|
options: dict[str, Any]
|
||||||
|
challenge: bytes
|
||||||
|
|
||||||
|
|
||||||
|
class VerifiedAuthentication(BaseModel):
|
||||||
|
"""
|
||||||
|
Verified WebAuthn authentication data
|
||||||
|
"""
|
||||||
|
new_sign_count: int
|
||||||
|
credential_id: bytes
|
||||||
|
user_verified: bool
|
||||||
|
backup_state: bool
|
||||||
|
backup_eligible: bool
|
||||||
|
|
||||||
|
|
||||||
|
class RegistrationBeginRequest(BaseModel):
|
||||||
|
"""
|
||||||
|
Request to begin passkey registration
|
||||||
|
"""
|
||||||
|
username: str = Field(
|
||||||
|
min_length = USERNAME_MIN_LENGTH,
|
||||||
|
max_length = USERNAME_MAX_LENGTH,
|
||||||
|
)
|
||||||
|
display_name: str = Field(
|
||||||
|
min_length = DISPLAY_NAME_MIN_LENGTH,
|
||||||
|
max_length = DISPLAY_NAME_MAX_LENGTH,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class RegistrationCompleteRequest(BaseModel):
|
||||||
|
"""
|
||||||
|
Request to complete passkey registration
|
||||||
|
"""
|
||||||
|
username: str = Field(min_length = USERNAME_MIN_LENGTH, max_length = USERNAME_MAX_LENGTH)
|
||||||
|
credential: dict[str, Any]
|
||||||
|
device_name: str | None = Field(
|
||||||
|
default = None,
|
||||||
|
max_length = DEVICE_NAME_MAX_LENGTH,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class AuthenticationBeginRequest(BaseModel):
|
||||||
|
"""
|
||||||
|
Request to begin passkey authentication
|
||||||
|
"""
|
||||||
|
username: str | None = Field(
|
||||||
|
default = None,
|
||||||
|
min_length = USERNAME_MIN_LENGTH,
|
||||||
|
max_length = USERNAME_MAX_LENGTH,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class AuthenticationCompleteRequest(BaseModel):
|
||||||
|
"""
|
||||||
|
Request to complete passkey authentication
|
||||||
|
"""
|
||||||
|
credential: dict[str, Any]
|
||||||
|
|
||||||
|
|
||||||
|
class UserResponse(BaseModel):
|
||||||
|
"""
|
||||||
|
User data response
|
||||||
|
"""
|
||||||
|
id: str
|
||||||
|
username: str
|
||||||
|
display_name: str
|
||||||
|
is_active: bool
|
||||||
|
is_verified: bool
|
||||||
|
created_at: str
|
||||||
|
|
||||||
|
|
||||||
|
class UserSearchRequest(BaseModel):
|
||||||
|
"""
|
||||||
|
Request to search for users
|
||||||
|
"""
|
||||||
|
query: str = Field(
|
||||||
|
min_length = USER_SEARCH_MIN_LENGTH,
|
||||||
|
max_length = USERNAME_MAX_LENGTH,
|
||||||
|
)
|
||||||
|
limit: int = Field(
|
||||||
|
default = USER_SEARCH_DEFAULT_LIMIT,
|
||||||
|
ge = 1,
|
||||||
|
le = USER_SEARCH_MAX_LIMIT,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class UserSearchResponse(BaseModel):
|
||||||
|
"""
|
||||||
|
Response containing search results
|
||||||
|
"""
|
||||||
|
users: list[UserResponse]
|
||||||
|
|
@ -0,0 +1,23 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Common Pydantic schemas for API responses
|
||||||
|
"""
|
||||||
|
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
|
||||||
|
class RootResponse(BaseModel):
|
||||||
|
"""
|
||||||
|
Root endpoint response schema
|
||||||
|
"""
|
||||||
|
app: str
|
||||||
|
version: str
|
||||||
|
status: str
|
||||||
|
environment: str
|
||||||
|
|
||||||
|
|
||||||
|
class HealthResponse(BaseModel):
|
||||||
|
"""
|
||||||
|
Health check endpoint response schema
|
||||||
|
"""
|
||||||
|
status: str
|
||||||
|
|
@ -0,0 +1,73 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Pydantic schemas for SurrealDB responses
|
||||||
|
"""
|
||||||
|
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
from app.core.enums import PresenceStatus, RoomType
|
||||||
|
|
||||||
|
|
||||||
|
class MessageResponse(BaseModel):
|
||||||
|
"""
|
||||||
|
Message response from SurrealDB
|
||||||
|
"""
|
||||||
|
id: str
|
||||||
|
room_id: str
|
||||||
|
sender_id: str
|
||||||
|
encrypted_content: str
|
||||||
|
encrypted_header: str
|
||||||
|
created_at: datetime
|
||||||
|
updated_at: datetime
|
||||||
|
|
||||||
|
|
||||||
|
class RoomResponse(BaseModel):
|
||||||
|
"""
|
||||||
|
Room response from SurrealDB
|
||||||
|
"""
|
||||||
|
id: str
|
||||||
|
name: str
|
||||||
|
room_type: RoomType
|
||||||
|
created_by: str
|
||||||
|
created_at: datetime
|
||||||
|
updated_at: datetime
|
||||||
|
is_ephemeral: bool = False
|
||||||
|
ttl_seconds: int | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class PresenceResponse(BaseModel):
|
||||||
|
"""
|
||||||
|
Presence response from SurrealDB
|
||||||
|
"""
|
||||||
|
id: str
|
||||||
|
user_id: str
|
||||||
|
room_id: str | None = None
|
||||||
|
status: PresenceStatus
|
||||||
|
last_seen: datetime
|
||||||
|
updated_at: datetime
|
||||||
|
|
||||||
|
|
||||||
|
class LiveQueryUpdate(BaseModel):
|
||||||
|
"""
|
||||||
|
Live query update notification from SurrealDB
|
||||||
|
"""
|
||||||
|
action: str
|
||||||
|
result: MessageResponse | PresenceResponse | RoomResponse
|
||||||
|
|
||||||
|
|
||||||
|
class LiveMessageUpdate(BaseModel):
|
||||||
|
"""
|
||||||
|
Live message update notification
|
||||||
|
"""
|
||||||
|
action: str
|
||||||
|
result: MessageResponse
|
||||||
|
|
||||||
|
|
||||||
|
class LivePresenceUpdate(BaseModel):
|
||||||
|
"""
|
||||||
|
Live presence update notification
|
||||||
|
"""
|
||||||
|
action: str
|
||||||
|
result: PresenceResponse
|
||||||
|
|
@ -0,0 +1,93 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Pydantic schemas for WebSocket message types
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Any
|
||||||
|
from datetime import datetime
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
ENCRYPTED_CONTENT_MAX_LENGTH,
|
||||||
|
MESSAGE_ID_MAX_LENGTH,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class BaseWSMessage(BaseModel):
|
||||||
|
"""
|
||||||
|
Base WebSocket message with common fields
|
||||||
|
"""
|
||||||
|
type: str
|
||||||
|
timestamp: datetime | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class EncryptedMessageWS(BaseWSMessage):
|
||||||
|
"""
|
||||||
|
Encrypted message sent over WebSocket
|
||||||
|
"""
|
||||||
|
type: str = "encrypted_message"
|
||||||
|
message_id: str = Field(max_length = MESSAGE_ID_MAX_LENGTH)
|
||||||
|
sender_id: str
|
||||||
|
recipient_id: str
|
||||||
|
ciphertext: str = Field(max_length = ENCRYPTED_CONTENT_MAX_LENGTH)
|
||||||
|
nonce: str
|
||||||
|
header: str
|
||||||
|
sender_username: str
|
||||||
|
|
||||||
|
|
||||||
|
class TypingIndicatorWS(BaseWSMessage):
|
||||||
|
"""
|
||||||
|
Typing indicator message
|
||||||
|
"""
|
||||||
|
type: str = "typing"
|
||||||
|
user_id: str
|
||||||
|
room_id: str
|
||||||
|
is_typing: bool
|
||||||
|
|
||||||
|
|
||||||
|
class PresenceUpdateWS(BaseWSMessage):
|
||||||
|
"""
|
||||||
|
User presence update message
|
||||||
|
"""
|
||||||
|
type: str = "presence"
|
||||||
|
user_id: str
|
||||||
|
status: str
|
||||||
|
last_seen: datetime
|
||||||
|
|
||||||
|
|
||||||
|
class ReadReceiptWS(BaseWSMessage):
|
||||||
|
"""
|
||||||
|
Message read receipt
|
||||||
|
"""
|
||||||
|
type: str = "receipt"
|
||||||
|
message_id: str = Field(max_length = MESSAGE_ID_MAX_LENGTH)
|
||||||
|
user_id: str
|
||||||
|
read_at: datetime
|
||||||
|
|
||||||
|
|
||||||
|
class ErrorMessageWS(BaseWSMessage):
|
||||||
|
"""
|
||||||
|
Error message sent over WebSocket
|
||||||
|
"""
|
||||||
|
type: str = "error"
|
||||||
|
error_code: str
|
||||||
|
error_message: str
|
||||||
|
details: dict[str, Any] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class WSConnectionRequest(BaseModel):
|
||||||
|
"""
|
||||||
|
WebSocket connection request with auth token
|
||||||
|
"""
|
||||||
|
user_id: UUID
|
||||||
|
token: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class WSHeartbeat(BaseModel):
|
||||||
|
"""
|
||||||
|
WebSocket heartbeat ping/pong message
|
||||||
|
"""
|
||||||
|
type: str = "heartbeat"
|
||||||
|
timestamp: datetime
|
||||||
|
|
@ -0,0 +1,24 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Service layer exports
|
||||||
|
"""
|
||||||
|
|
||||||
|
from app.services.auth_service import AuthService, auth_service
|
||||||
|
from app.services.message_service import MessageService, message_service
|
||||||
|
from app.services.prekey_service import PrekeyService, prekey_service
|
||||||
|
from app.services.presence_service import PresenceService, presence_service
|
||||||
|
from app.services.websocket_service import WebSocketService, websocket_service
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"AuthService",
|
||||||
|
"auth_service",
|
||||||
|
"MessageService",
|
||||||
|
"message_service",
|
||||||
|
"PrekeyService",
|
||||||
|
"prekey_service",
|
||||||
|
"PresenceService",
|
||||||
|
"presence_service",
|
||||||
|
"WebSocketService",
|
||||||
|
"websocket_service",
|
||||||
|
]
|
||||||
|
|
@ -0,0 +1,580 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Authentication service for user and credential management
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from typing import Any
|
||||||
|
from uuid import UUID
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
|
||||||
|
from sqlmodel import select
|
||||||
|
from sqlalchemy.exc import IntegrityError
|
||||||
|
from sqlalchemy.orm import selectinload
|
||||||
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||||
|
from webauthn.helpers import base64url_to_bytes, bytes_to_base64url
|
||||||
|
|
||||||
|
from app.config import USER_SEARCH_DEFAULT_LIMIT
|
||||||
|
from app.core.exceptions import (
|
||||||
|
ChallengeExpiredError,
|
||||||
|
CredentialNotFoundError,
|
||||||
|
CredentialVerificationError,
|
||||||
|
DatabaseError,
|
||||||
|
InvalidDataError,
|
||||||
|
UserExistsError,
|
||||||
|
UserInactiveError,
|
||||||
|
UserNotFoundError,
|
||||||
|
)
|
||||||
|
from app.core.passkey.passkey_manager import passkey_manager
|
||||||
|
from app.core.redis_manager import redis_manager
|
||||||
|
from app.models.Credential import Credential
|
||||||
|
from app.models.User import User
|
||||||
|
from app.schemas.auth import (
|
||||||
|
AuthenticationBeginRequest,
|
||||||
|
AuthenticationCompleteRequest,
|
||||||
|
RegistrationBeginRequest,
|
||||||
|
RegistrationCompleteRequest,
|
||||||
|
UserResponse,
|
||||||
|
VerifiedRegistration,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class AuthService:
|
||||||
|
"""
|
||||||
|
Service for managing user authentication and credentials
|
||||||
|
"""
|
||||||
|
async def create_user(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
username: str,
|
||||||
|
display_name: str,
|
||||||
|
) -> User:
|
||||||
|
"""
|
||||||
|
Create a new user with username uniqueness check
|
||||||
|
"""
|
||||||
|
statement = select(User).where(User.username == username)
|
||||||
|
result = await session.execute(statement)
|
||||||
|
existing_user = result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if existing_user:
|
||||||
|
logger.warning("Attempted to create duplicate user: %s", username)
|
||||||
|
raise UserExistsError(f"Username {username} already exists")
|
||||||
|
|
||||||
|
user = User(
|
||||||
|
username = username,
|
||||||
|
display_name = display_name,
|
||||||
|
)
|
||||||
|
|
||||||
|
session.add(user)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
await session.refresh(user)
|
||||||
|
logger.info("Created new user: %s (ID: %s)", username, user.id)
|
||||||
|
return user
|
||||||
|
except IntegrityError as e:
|
||||||
|
await session.rollback()
|
||||||
|
logger.error(
|
||||||
|
"Database integrity error creating user %s: %s",
|
||||||
|
username,
|
||||||
|
e
|
||||||
|
)
|
||||||
|
raise DatabaseError(
|
||||||
|
"Failed to create user: database constraint violation"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
async def store_credential(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
user_id: UUID,
|
||||||
|
verified: VerifiedRegistration,
|
||||||
|
device_name: str | None = None,
|
||||||
|
) -> Credential:
|
||||||
|
"""
|
||||||
|
Store WebAuthn credential after successful registration
|
||||||
|
"""
|
||||||
|
credential = Credential(
|
||||||
|
user_id = user_id,
|
||||||
|
credential_id = bytes_to_base64url(verified.credential_id),
|
||||||
|
public_key = bytes_to_base64url(verified.credential_public_key),
|
||||||
|
sign_count = verified.sign_count,
|
||||||
|
aaguid = bytes_to_base64url(verified.aaguid),
|
||||||
|
backup_eligible = verified.backup_eligible,
|
||||||
|
backup_state = verified.backup_state,
|
||||||
|
attestation_type = verified.attestation_format,
|
||||||
|
device_name = device_name,
|
||||||
|
last_used_at = datetime.now(UTC),
|
||||||
|
)
|
||||||
|
|
||||||
|
session.add(credential)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
await session.refresh(credential)
|
||||||
|
logger.info(
|
||||||
|
"Stored credential %s... for user %s",
|
||||||
|
credential.credential_id[: 16],
|
||||||
|
user_id
|
||||||
|
)
|
||||||
|
return credential
|
||||||
|
except IntegrityError as e:
|
||||||
|
await session.rollback()
|
||||||
|
logger.error("Database integrity error storing credential: %s", e)
|
||||||
|
raise DatabaseError(
|
||||||
|
"Failed to store credential: database constraint violation"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
async def get_user_by_username(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
username: str,
|
||||||
|
) -> User | None:
|
||||||
|
"""
|
||||||
|
Retrieve user by username with credentials relationship eager loaded
|
||||||
|
"""
|
||||||
|
statement = (
|
||||||
|
select(User).where(User.username == username).options(
|
||||||
|
selectinload(User.credentials)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
result = await session.execute(statement)
|
||||||
|
user = result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if user:
|
||||||
|
logger.debug(
|
||||||
|
"Retrieved user %s with %s credentials",
|
||||||
|
username,
|
||||||
|
len(user.credentials)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
logger.debug("User not found: %s", username)
|
||||||
|
|
||||||
|
return user
|
||||||
|
|
||||||
|
async def get_user_by_id(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
user_id: UUID,
|
||||||
|
) -> User | None:
|
||||||
|
"""
|
||||||
|
Retrieve user by ID with credentials relationship eager loaded
|
||||||
|
"""
|
||||||
|
statement = (
|
||||||
|
select(User).where(User.id == user_id).options(
|
||||||
|
selectinload(User.credentials)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
result = await session.execute(statement)
|
||||||
|
user = result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if user:
|
||||||
|
logger.debug(
|
||||||
|
"Retrieved user %s with %s credentials",
|
||||||
|
user_id,
|
||||||
|
len(user.credentials)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
logger.debug("User not found: %s", user_id)
|
||||||
|
|
||||||
|
return user
|
||||||
|
|
||||||
|
async def search_users(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
query: str,
|
||||||
|
limit: int = USER_SEARCH_DEFAULT_LIMIT,
|
||||||
|
exclude_user_id: UUID | None = None,
|
||||||
|
) -> list[User]:
|
||||||
|
"""
|
||||||
|
Search for active users by username or display name
|
||||||
|
"""
|
||||||
|
search_pattern = f"%{query.lower()}%"
|
||||||
|
|
||||||
|
statement = (
|
||||||
|
select(User)
|
||||||
|
.where(
|
||||||
|
User.is_active == True,
|
||||||
|
(
|
||||||
|
User.username.ilike(search_pattern) |
|
||||||
|
User.display_name.ilike(search_pattern)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
.limit(limit)
|
||||||
|
)
|
||||||
|
|
||||||
|
if exclude_user_id is not None:
|
||||||
|
statement = statement.where(User.id != exclude_user_id)
|
||||||
|
|
||||||
|
result = await session.execute(statement)
|
||||||
|
users = result.scalars().all()
|
||||||
|
|
||||||
|
logger.debug(
|
||||||
|
"Search for '%s' returned %d users",
|
||||||
|
query,
|
||||||
|
len(users)
|
||||||
|
)
|
||||||
|
|
||||||
|
return list(users)
|
||||||
|
|
||||||
|
async def get_credential_by_id(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
credential_id: str,
|
||||||
|
) -> Credential | None:
|
||||||
|
"""
|
||||||
|
Retrieve credential by credential_id
|
||||||
|
"""
|
||||||
|
statement = select(Credential).where(
|
||||||
|
Credential.credential_id == credential_id
|
||||||
|
)
|
||||||
|
result = await session.execute(statement)
|
||||||
|
credential = result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if credential:
|
||||||
|
logger.debug("Retrieved credential %s...", credential_id[: 16])
|
||||||
|
else:
|
||||||
|
logger.debug("Credential not found: %s...", credential_id[: 16])
|
||||||
|
|
||||||
|
return credential
|
||||||
|
|
||||||
|
async def update_credential_counter(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
credential_id: str,
|
||||||
|
new_count: int,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Update credential signature counter after successful authentication
|
||||||
|
"""
|
||||||
|
statement = select(Credential).where(
|
||||||
|
Credential.credential_id == credential_id
|
||||||
|
)
|
||||||
|
result = await session.execute(statement)
|
||||||
|
credential = result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if not credential:
|
||||||
|
logger.error(
|
||||||
|
"Credential not found for counter update: %s...",
|
||||||
|
credential_id[: 16]
|
||||||
|
)
|
||||||
|
raise CredentialNotFoundError("Credential not found")
|
||||||
|
|
||||||
|
old_count = credential.sign_count
|
||||||
|
credential.sign_count = new_count
|
||||||
|
credential.last_used_at = datetime.now(UTC)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
logger.info(
|
||||||
|
"Updated credential %s... counter: %s -> %s",
|
||||||
|
credential_id[: 16],
|
||||||
|
old_count,
|
||||||
|
new_count
|
||||||
|
)
|
||||||
|
except IntegrityError as e:
|
||||||
|
await session.rollback()
|
||||||
|
logger.error("Database error updating credential counter: %s", e)
|
||||||
|
raise DatabaseError("Failed to update credential counter") from e
|
||||||
|
|
||||||
|
async def update_backup_state(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
credential_id: str,
|
||||||
|
backup_state: bool,
|
||||||
|
backup_eligible: bool,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Update credential backup flags (WebAuthn Level 3)
|
||||||
|
"""
|
||||||
|
statement = select(Credential).where(
|
||||||
|
Credential.credential_id == credential_id
|
||||||
|
)
|
||||||
|
result = await session.execute(statement)
|
||||||
|
credential = result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if not credential:
|
||||||
|
logger.error(
|
||||||
|
"Credential not found for backup state update: %s...",
|
||||||
|
credential_id[: 16]
|
||||||
|
)
|
||||||
|
raise CredentialNotFoundError("Credential not found")
|
||||||
|
|
||||||
|
if credential.backup_state != backup_state:
|
||||||
|
logger.warning(
|
||||||
|
"Credential %s... backup state changed: %s -> %s",
|
||||||
|
credential_id[: 16],
|
||||||
|
credential.backup_state,
|
||||||
|
backup_state
|
||||||
|
)
|
||||||
|
|
||||||
|
credential.backup_state = backup_state
|
||||||
|
credential.backup_eligible = backup_eligible
|
||||||
|
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
logger.debug(
|
||||||
|
"Updated credential %s... backup_state=%s, backup_eligible=%s",
|
||||||
|
credential_id[: 16],
|
||||||
|
backup_state,
|
||||||
|
backup_eligible
|
||||||
|
)
|
||||||
|
except IntegrityError as e:
|
||||||
|
await session.rollback()
|
||||||
|
logger.error("Database error updating backup state: %s", e)
|
||||||
|
raise DatabaseError("Failed to update backup state") from e
|
||||||
|
|
||||||
|
async def begin_registration(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
request: RegistrationBeginRequest,
|
||||||
|
) -> dict[str,
|
||||||
|
Any]:
|
||||||
|
"""
|
||||||
|
Begin WebAuthn passkey registration flow
|
||||||
|
"""
|
||||||
|
existing_user = await self.get_user_by_username(
|
||||||
|
session = session,
|
||||||
|
username = request.username,
|
||||||
|
)
|
||||||
|
|
||||||
|
if existing_user:
|
||||||
|
logger.warning(
|
||||||
|
"Registration attempt for existing user: %s",
|
||||||
|
request.username
|
||||||
|
)
|
||||||
|
raise UserExistsError(f"Username {request.username} already exists")
|
||||||
|
|
||||||
|
user_id_bytes = request.username.encode()
|
||||||
|
|
||||||
|
exclude_credentials = []
|
||||||
|
if existing_user:
|
||||||
|
exclude_credentials = [
|
||||||
|
base64url_to_bytes(cred.credential_id)
|
||||||
|
for cred in existing_user.credentials
|
||||||
|
]
|
||||||
|
|
||||||
|
options_response = passkey_manager.generate_registration_options(
|
||||||
|
user_id = user_id_bytes,
|
||||||
|
username = request.username,
|
||||||
|
display_name = request.display_name,
|
||||||
|
exclude_credentials = exclude_credentials,
|
||||||
|
)
|
||||||
|
|
||||||
|
await redis_manager.set_registration_challenge(
|
||||||
|
user_id = request.username,
|
||||||
|
challenge = options_response.challenge,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info("Started registration for user: %s", request.username)
|
||||||
|
return options_response.options
|
||||||
|
|
||||||
|
async def complete_registration(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
request: RegistrationCompleteRequest,
|
||||||
|
username: str,
|
||||||
|
) -> UserResponse:
|
||||||
|
"""
|
||||||
|
Complete WebAuthn passkey registration
|
||||||
|
"""
|
||||||
|
expected_challenge = await redis_manager.get_registration_challenge(
|
||||||
|
user_id = username
|
||||||
|
)
|
||||||
|
|
||||||
|
if not expected_challenge:
|
||||||
|
logger.warning(
|
||||||
|
"Registration challenge not found or expired for user: %s",
|
||||||
|
username
|
||||||
|
)
|
||||||
|
raise ChallengeExpiredError(
|
||||||
|
"Challenge expired or not found - please restart registration"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
verified = passkey_manager.verify_registration(
|
||||||
|
credential = request.credential,
|
||||||
|
expected_challenge = expected_challenge,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Registration verification failed: %s", e)
|
||||||
|
raise CredentialVerificationError(
|
||||||
|
f"Registration verification failed: {str(e)}"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
user = await self.create_user(
|
||||||
|
session = session,
|
||||||
|
username = username,
|
||||||
|
display_name = request.credential.get("displayName",
|
||||||
|
username),
|
||||||
|
)
|
||||||
|
|
||||||
|
await self.store_credential(
|
||||||
|
session = session,
|
||||||
|
user_id = user.id,
|
||||||
|
verified = verified,
|
||||||
|
device_name = request.device_name,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info("Registration completed for user: %s", username)
|
||||||
|
|
||||||
|
return UserResponse(
|
||||||
|
id = str(user.id),
|
||||||
|
username = user.username,
|
||||||
|
display_name = user.display_name,
|
||||||
|
is_active = user.is_active,
|
||||||
|
is_verified = user.is_verified,
|
||||||
|
created_at = user.created_at.isoformat(),
|
||||||
|
)
|
||||||
|
|
||||||
|
async def begin_authentication(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
request: AuthenticationBeginRequest,
|
||||||
|
) -> dict[str,
|
||||||
|
Any]:
|
||||||
|
"""
|
||||||
|
Begin WebAuthn passkey authentication flow
|
||||||
|
"""
|
||||||
|
allow_credentials = None
|
||||||
|
|
||||||
|
if request.username:
|
||||||
|
user = await self.get_user_by_username(
|
||||||
|
session = session,
|
||||||
|
username = request.username,
|
||||||
|
)
|
||||||
|
|
||||||
|
if not user:
|
||||||
|
logger.warning(
|
||||||
|
"Authentication attempt for non-existent user: %s",
|
||||||
|
request.username
|
||||||
|
)
|
||||||
|
raise UserNotFoundError("User not found")
|
||||||
|
|
||||||
|
if not user.is_active:
|
||||||
|
logger.warning(
|
||||||
|
"Authentication attempt for inactive user: %s",
|
||||||
|
request.username
|
||||||
|
)
|
||||||
|
raise UserInactiveError("User account is inactive")
|
||||||
|
|
||||||
|
allow_credentials = [
|
||||||
|
base64url_to_bytes(cred.credential_id)
|
||||||
|
for cred in user.credentials
|
||||||
|
]
|
||||||
|
|
||||||
|
options_response = passkey_manager.generate_authentication_options(
|
||||||
|
allow_credentials = allow_credentials,
|
||||||
|
)
|
||||||
|
|
||||||
|
user_id = request.username if request.username else "discoverable"
|
||||||
|
await redis_manager.set_authentication_challenge(
|
||||||
|
user_id = user_id,
|
||||||
|
challenge = options_response.challenge,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info("Started authentication for user: %s", user_id)
|
||||||
|
return options_response.options
|
||||||
|
|
||||||
|
async def complete_authentication(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
request: AuthenticationCompleteRequest,
|
||||||
|
) -> UserResponse:
|
||||||
|
"""
|
||||||
|
Complete WebAuthn passkey authentication
|
||||||
|
"""
|
||||||
|
credential_id = request.credential.get("id")
|
||||||
|
if not credential_id:
|
||||||
|
raise InvalidDataError("Missing credential ID")
|
||||||
|
|
||||||
|
credential = await self.get_credential_by_id(
|
||||||
|
session = session,
|
||||||
|
credential_id = credential_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
if not credential:
|
||||||
|
logger.warning(
|
||||||
|
"Authentication with unknown credential: %s...",
|
||||||
|
credential_id[: 16]
|
||||||
|
)
|
||||||
|
raise CredentialNotFoundError("Credential not found")
|
||||||
|
|
||||||
|
user = await self.get_user_by_id(
|
||||||
|
session = session,
|
||||||
|
user_id = credential.user_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
if not user:
|
||||||
|
logger.error(
|
||||||
|
"User not found for credential: %s...",
|
||||||
|
credential_id[: 16]
|
||||||
|
)
|
||||||
|
raise UserNotFoundError("User not found")
|
||||||
|
|
||||||
|
if not user.is_active:
|
||||||
|
logger.warning(
|
||||||
|
"Authentication attempt for inactive user: %s",
|
||||||
|
user.username
|
||||||
|
)
|
||||||
|
raise UserInactiveError("User account is inactive")
|
||||||
|
|
||||||
|
expected_challenge = await redis_manager.get_authentication_challenge(
|
||||||
|
user_id = user.username
|
||||||
|
)
|
||||||
|
|
||||||
|
if not expected_challenge:
|
||||||
|
logger.warning(
|
||||||
|
"Authentication challenge not found for user: %s",
|
||||||
|
user.username
|
||||||
|
)
|
||||||
|
raise ChallengeExpiredError(
|
||||||
|
"Challenge expired or not found - please restart authentication"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
verified = passkey_manager.verify_authentication(
|
||||||
|
credential = request.credential,
|
||||||
|
expected_challenge = expected_challenge,
|
||||||
|
credential_public_key = base64url_to_bytes(credential.public_key),
|
||||||
|
credential_current_sign_count = credential.sign_count,
|
||||||
|
)
|
||||||
|
except ValueError as e:
|
||||||
|
logger.error("Authentication verification failed: %s", e)
|
||||||
|
raise CredentialVerificationError(str(e)) from e
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Unexpected error during authentication: %s", e)
|
||||||
|
raise CredentialVerificationError(
|
||||||
|
"Authentication verification failed"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
await self.update_credential_counter(
|
||||||
|
session = session,
|
||||||
|
credential_id = credential.credential_id,
|
||||||
|
new_count = verified.new_sign_count,
|
||||||
|
)
|
||||||
|
|
||||||
|
if (credential.backup_state != verified.backup_state
|
||||||
|
or credential.backup_eligible != verified.backup_eligible):
|
||||||
|
await self.update_backup_state(
|
||||||
|
session = session,
|
||||||
|
credential_id = credential.credential_id,
|
||||||
|
backup_state = verified.backup_state,
|
||||||
|
backup_eligible = verified.backup_eligible,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info("Authentication successful for user: %s", user.username)
|
||||||
|
|
||||||
|
return UserResponse(
|
||||||
|
id = str(user.id),
|
||||||
|
username = user.username,
|
||||||
|
display_name = user.display_name,
|
||||||
|
is_active = user.is_active,
|
||||||
|
is_verified = user.is_verified,
|
||||||
|
created_at = user.created_at.isoformat(),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
auth_service = AuthService()
|
||||||
|
|
@ -0,0 +1,414 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Message service with end-to-end encryption using Double Ratchet
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from typing import Any
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from sqlmodel import select
|
||||||
|
from sqlalchemy.exc import IntegrityError
|
||||||
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||||
|
from cryptography.hazmat.primitives import serialization
|
||||||
|
from webauthn.helpers import base64url_to_bytes, bytes_to_base64url
|
||||||
|
from cryptography.hazmat.primitives.asymmetric.x25519 import (
|
||||||
|
X25519PrivateKey,
|
||||||
|
)
|
||||||
|
from app.core.encryption.double_ratchet import (
|
||||||
|
DoubleRatchetState,
|
||||||
|
EncryptedMessage,
|
||||||
|
double_ratchet,
|
||||||
|
)
|
||||||
|
from app.core.encryption.x3dh_manager import x3dh_manager
|
||||||
|
from app.core.exceptions import (
|
||||||
|
DatabaseError,
|
||||||
|
DecryptionError,
|
||||||
|
EncryptionError,
|
||||||
|
InvalidDataError,
|
||||||
|
KeyExchangeError,
|
||||||
|
RatchetStateNotFoundError,
|
||||||
|
UserNotFoundError,
|
||||||
|
)
|
||||||
|
from app.core.surreal_manager import surreal_db
|
||||||
|
from app.models.IdentityKey import IdentityKey
|
||||||
|
from app.models.RatchetState import RatchetState
|
||||||
|
from app.models.User import User
|
||||||
|
from app.services.prekey_service import prekey_service
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class MessageService:
|
||||||
|
"""
|
||||||
|
Service for encrypted messaging using Double Ratchet protocol
|
||||||
|
"""
|
||||||
|
async def initialize_conversation(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
sender_id: UUID,
|
||||||
|
recipient_id: UUID
|
||||||
|
) -> RatchetState:
|
||||||
|
"""
|
||||||
|
Performs X3DH key exchange and initializes Double Ratchet for new conversation
|
||||||
|
"""
|
||||||
|
if sender_id == recipient_id:
|
||||||
|
raise InvalidDataError("Cannot start conversation with yourself")
|
||||||
|
|
||||||
|
existing_state_statement = select(RatchetState).where(
|
||||||
|
RatchetState.user_id == sender_id,
|
||||||
|
RatchetState.peer_user_id == recipient_id
|
||||||
|
)
|
||||||
|
existing_state_result = await session.execute(existing_state_statement)
|
||||||
|
existing_state = existing_state_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if existing_state:
|
||||||
|
logger.warning(
|
||||||
|
"Ratchet state already exists for %s -> %s",
|
||||||
|
sender_id,
|
||||||
|
recipient_id
|
||||||
|
)
|
||||||
|
return existing_state
|
||||||
|
|
||||||
|
sender_ik_statement = select(IdentityKey).where(
|
||||||
|
IdentityKey.user_id == sender_id
|
||||||
|
)
|
||||||
|
sender_ik_result = await session.execute(sender_ik_statement)
|
||||||
|
sender_ik = sender_ik_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if not sender_ik:
|
||||||
|
logger.error("Sender identity key not found: %s", sender_id)
|
||||||
|
raise InvalidDataError(
|
||||||
|
"Sender has no identity key - initialize encryption first"
|
||||||
|
)
|
||||||
|
|
||||||
|
recipient_bundle = await prekey_service.get_prekey_bundle(
|
||||||
|
session,
|
||||||
|
recipient_id
|
||||||
|
)
|
||||||
|
|
||||||
|
recipient_ik_statement = select(IdentityKey).where(
|
||||||
|
IdentityKey.user_id == recipient_id
|
||||||
|
)
|
||||||
|
recipient_ik_result = await session.execute(recipient_ik_statement)
|
||||||
|
recipient_ik = recipient_ik_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if not recipient_ik:
|
||||||
|
logger.error("Recipient identity key not found: %s", recipient_id)
|
||||||
|
raise InvalidDataError("Recipient has no identity key")
|
||||||
|
|
||||||
|
try:
|
||||||
|
x3dh_result = x3dh_manager.perform_x3dh_sender(
|
||||||
|
alice_identity_private_x25519 = sender_ik.private_key,
|
||||||
|
bob_bundle = recipient_bundle,
|
||||||
|
bob_identity_public_ed25519 = recipient_ik.public_key_ed25519
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("X3DH key exchange failed: %s", e)
|
||||||
|
raise KeyExchangeError(f"Key exchange failed: {str(e)}") from e
|
||||||
|
|
||||||
|
recipient_spk_public_bytes = base64url_to_bytes(
|
||||||
|
recipient_bundle.signed_prekey
|
||||||
|
)
|
||||||
|
|
||||||
|
dr_state = double_ratchet.initialize_sender(
|
||||||
|
shared_key = x3dh_result.shared_key,
|
||||||
|
peer_public_key = recipient_spk_public_bytes
|
||||||
|
)
|
||||||
|
|
||||||
|
dh_private_bytes = dr_state.dh_private_key.private_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PrivateFormat.Raw,
|
||||||
|
encryption_algorithm = serialization.NoEncryption()
|
||||||
|
) if dr_state.dh_private_key else b''
|
||||||
|
|
||||||
|
dh_public_bytes = dr_state.dh_private_key.public_key().public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
) if dr_state.dh_private_key else b''
|
||||||
|
|
||||||
|
ratchet_state = RatchetState(
|
||||||
|
user_id = sender_id,
|
||||||
|
peer_user_id = recipient_id,
|
||||||
|
dh_private_key = bytes_to_base64url(dh_private_bytes),
|
||||||
|
dh_public_key = bytes_to_base64url(dh_public_bytes),
|
||||||
|
dh_peer_public_key = bytes_to_base64url(dr_state.dh_peer_public_key)
|
||||||
|
if dr_state.dh_peer_public_key else None,
|
||||||
|
root_key = bytes_to_base64url(dr_state.root_key),
|
||||||
|
sending_chain_key = bytes_to_base64url(dr_state.sending_chain_key),
|
||||||
|
receiving_chain_key = bytes_to_base64url(
|
||||||
|
dr_state.receiving_chain_key
|
||||||
|
),
|
||||||
|
sending_message_number = dr_state.sending_message_number,
|
||||||
|
receiving_message_number = dr_state.receiving_message_number,
|
||||||
|
previous_sending_chain_length = (
|
||||||
|
dr_state.previous_sending_chain_length
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
session.add(ratchet_state)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
await session.refresh(ratchet_state)
|
||||||
|
logger.info(
|
||||||
|
"Initialized conversation: %s -> %s (X3DH complete)",
|
||||||
|
sender_id,
|
||||||
|
recipient_id
|
||||||
|
)
|
||||||
|
except IntegrityError as e:
|
||||||
|
await session.rollback()
|
||||||
|
logger.error("Database error saving ratchet state: %s", e)
|
||||||
|
raise DatabaseError("Failed to initialize conversation") from e
|
||||||
|
|
||||||
|
return ratchet_state
|
||||||
|
|
||||||
|
async def _load_ratchet_state_from_db(
|
||||||
|
self,
|
||||||
|
ratchet_state_db: RatchetState
|
||||||
|
) -> DoubleRatchetState:
|
||||||
|
"""
|
||||||
|
Converts database RatchetState to DoubleRatchetState object
|
||||||
|
"""
|
||||||
|
dh_private_key = None
|
||||||
|
if ratchet_state_db.dh_private_key:
|
||||||
|
dh_private_bytes = base64url_to_bytes(ratchet_state_db.dh_private_key)
|
||||||
|
dh_private_key = X25519PrivateKey.from_private_bytes(dh_private_bytes)
|
||||||
|
|
||||||
|
dh_peer_public_key = None
|
||||||
|
if ratchet_state_db.dh_peer_public_key:
|
||||||
|
dh_peer_public_key = base64url_to_bytes(
|
||||||
|
ratchet_state_db.dh_peer_public_key
|
||||||
|
)
|
||||||
|
|
||||||
|
root_key = base64url_to_bytes(ratchet_state_db.root_key)
|
||||||
|
sending_chain_key = base64url_to_bytes(ratchet_state_db.sending_chain_key)
|
||||||
|
receiving_chain_key = base64url_to_bytes(
|
||||||
|
ratchet_state_db.receiving_chain_key
|
||||||
|
)
|
||||||
|
|
||||||
|
return DoubleRatchetState(
|
||||||
|
root_key = root_key,
|
||||||
|
sending_chain_key = sending_chain_key,
|
||||||
|
receiving_chain_key = receiving_chain_key,
|
||||||
|
dh_private_key = dh_private_key,
|
||||||
|
dh_peer_public_key = dh_peer_public_key,
|
||||||
|
sending_message_number = ratchet_state_db.sending_message_number,
|
||||||
|
receiving_message_number = (
|
||||||
|
ratchet_state_db.receiving_message_number
|
||||||
|
),
|
||||||
|
previous_sending_chain_length = (
|
||||||
|
ratchet_state_db.previous_sending_chain_length
|
||||||
|
),
|
||||||
|
skipped_message_keys = {}
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _save_ratchet_state_to_db(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
ratchet_state_db: RatchetState,
|
||||||
|
dr_state: DoubleRatchetState
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Updates database RatchetState from DoubleRatchetState object
|
||||||
|
"""
|
||||||
|
if dr_state.dh_private_key:
|
||||||
|
dh_private_bytes = dr_state.dh_private_key.private_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PrivateFormat.Raw,
|
||||||
|
encryption_algorithm = serialization.NoEncryption()
|
||||||
|
)
|
||||||
|
dh_public_bytes = dr_state.dh_private_key.public_key().public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
ratchet_state_db.dh_private_key = bytes_to_base64url(dh_private_bytes)
|
||||||
|
ratchet_state_db.dh_public_key = bytes_to_base64url(dh_public_bytes)
|
||||||
|
else:
|
||||||
|
ratchet_state_db.dh_private_key = None
|
||||||
|
ratchet_state_db.dh_public_key = None
|
||||||
|
|
||||||
|
if dr_state.dh_peer_public_key:
|
||||||
|
ratchet_state_db.dh_peer_public_key = bytes_to_base64url(
|
||||||
|
dr_state.dh_peer_public_key
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
ratchet_state_db.dh_peer_public_key = None
|
||||||
|
|
||||||
|
ratchet_state_db.root_key = bytes_to_base64url(dr_state.root_key)
|
||||||
|
ratchet_state_db.sending_chain_key = bytes_to_base64url(
|
||||||
|
dr_state.sending_chain_key
|
||||||
|
)
|
||||||
|
ratchet_state_db.receiving_chain_key = bytes_to_base64url(
|
||||||
|
dr_state.receiving_chain_key
|
||||||
|
)
|
||||||
|
ratchet_state_db.sending_message_number = (
|
||||||
|
dr_state.sending_message_number
|
||||||
|
)
|
||||||
|
ratchet_state_db.receiving_message_number = (
|
||||||
|
dr_state.receiving_message_number
|
||||||
|
)
|
||||||
|
ratchet_state_db.previous_sending_chain_length = (
|
||||||
|
dr_state.previous_sending_chain_length
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
logger.debug(
|
||||||
|
"Saved ratchet state: send=%s, recv=%s",
|
||||||
|
dr_state.sending_message_number,
|
||||||
|
dr_state.receiving_message_number
|
||||||
|
)
|
||||||
|
except IntegrityError as e:
|
||||||
|
await session.rollback()
|
||||||
|
logger.error("Database error saving ratchet state: %s", e)
|
||||||
|
raise DatabaseError("Failed to save ratchet state") from e
|
||||||
|
|
||||||
|
async def send_encrypted_message(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
sender_id: UUID,
|
||||||
|
recipient_id: UUID,
|
||||||
|
plaintext: str
|
||||||
|
) -> Any:
|
||||||
|
"""
|
||||||
|
Encrypts message with Double Ratchet and stores in SurrealDB
|
||||||
|
"""
|
||||||
|
ratchet_state_statement = select(RatchetState).where(
|
||||||
|
RatchetState.user_id == sender_id,
|
||||||
|
RatchetState.peer_user_id == recipient_id
|
||||||
|
)
|
||||||
|
ratchet_state_result = await session.execute(ratchet_state_statement)
|
||||||
|
ratchet_state_db = ratchet_state_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if not ratchet_state_db:
|
||||||
|
logger.warning(
|
||||||
|
"No ratchet state for %s -> %s, initializing",
|
||||||
|
sender_id,
|
||||||
|
recipient_id
|
||||||
|
)
|
||||||
|
ratchet_state_db = await self.initialize_conversation(
|
||||||
|
session,
|
||||||
|
sender_id,
|
||||||
|
recipient_id
|
||||||
|
)
|
||||||
|
|
||||||
|
dr_state = await self._load_ratchet_state_from_db(ratchet_state_db)
|
||||||
|
|
||||||
|
sender_user_statement = select(User).where(User.id == sender_id)
|
||||||
|
sender_user_result = await session.execute(sender_user_statement)
|
||||||
|
sender_user = sender_user_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if not sender_user:
|
||||||
|
raise UserNotFoundError("Sender not found")
|
||||||
|
|
||||||
|
associated_data = f"{sender_id}:{recipient_id}".encode()
|
||||||
|
|
||||||
|
try:
|
||||||
|
encrypted_msg = double_ratchet.encrypt_message(
|
||||||
|
dr_state,
|
||||||
|
plaintext.encode(),
|
||||||
|
associated_data
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Encryption failed: %s", e)
|
||||||
|
raise EncryptionError(f"Failed to encrypt message: {str(e)}") from e
|
||||||
|
|
||||||
|
await self._save_ratchet_state_to_db(session, ratchet_state_db, dr_state)
|
||||||
|
|
||||||
|
message_header = {
|
||||||
|
"dh_public_key": bytes_to_base64url(encrypted_msg.dh_public_key),
|
||||||
|
"message_number": encrypted_msg.message_number,
|
||||||
|
"previous_chain_length": encrypted_msg.previous_chain_length
|
||||||
|
}
|
||||||
|
|
||||||
|
surreal_message = {
|
||||||
|
"sender_id": str(sender_id),
|
||||||
|
"recipient_id": str(recipient_id),
|
||||||
|
"ciphertext": bytes_to_base64url(encrypted_msg.ciphertext),
|
||||||
|
"nonce": bytes_to_base64url(encrypted_msg.nonce),
|
||||||
|
"header": json.dumps(message_header),
|
||||||
|
"sender_username": sender_user.username
|
||||||
|
}
|
||||||
|
|
||||||
|
try:
|
||||||
|
result = await surreal_db.create_message(surreal_message)
|
||||||
|
logger.info(
|
||||||
|
"Sent encrypted message: %s -> %s (msg #%s)",
|
||||||
|
sender_id,
|
||||||
|
recipient_id,
|
||||||
|
encrypted_msg.message_number
|
||||||
|
)
|
||||||
|
return result
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to store encrypted message: %s", e)
|
||||||
|
raise DatabaseError(f"Failed to store message: {str(e)}") from e
|
||||||
|
|
||||||
|
async def decrypt_received_message(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
recipient_id: UUID,
|
||||||
|
message_data: dict[str,
|
||||||
|
Any]
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Decrypts received message using Double Ratchet
|
||||||
|
"""
|
||||||
|
sender_id = UUID(message_data["sender_id"])
|
||||||
|
|
||||||
|
ratchet_state_statement = select(RatchetState).where(
|
||||||
|
RatchetState.user_id == recipient_id,
|
||||||
|
RatchetState.peer_user_id == sender_id
|
||||||
|
)
|
||||||
|
ratchet_state_result = await session.execute(ratchet_state_statement)
|
||||||
|
ratchet_state_db = ratchet_state_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if not ratchet_state_db:
|
||||||
|
logger.error(
|
||||||
|
"No ratchet state for receiving: %s <- %s",
|
||||||
|
recipient_id,
|
||||||
|
sender_id
|
||||||
|
)
|
||||||
|
raise RatchetStateNotFoundError("No encryption session with sender")
|
||||||
|
|
||||||
|
dr_state = await self._load_ratchet_state_from_db(ratchet_state_db)
|
||||||
|
|
||||||
|
header = json.loads(message_data["header"])
|
||||||
|
|
||||||
|
encrypted_msg = EncryptedMessage(
|
||||||
|
ciphertext = base64url_to_bytes(message_data["ciphertext"]),
|
||||||
|
nonce = base64url_to_bytes(message_data["nonce"]),
|
||||||
|
dh_public_key = base64url_to_bytes(header["dh_public_key"]),
|
||||||
|
message_number = header["message_number"],
|
||||||
|
previous_chain_length = header["previous_chain_length"]
|
||||||
|
)
|
||||||
|
|
||||||
|
associated_data = f"{sender_id}:{recipient_id}".encode()
|
||||||
|
|
||||||
|
try:
|
||||||
|
plaintext_bytes = double_ratchet.decrypt_message(
|
||||||
|
dr_state,
|
||||||
|
encrypted_msg,
|
||||||
|
associated_data
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Decryption failed: %s", e)
|
||||||
|
raise DecryptionError(f"Failed to decrypt message: {str(e)}") from e
|
||||||
|
|
||||||
|
await self._save_ratchet_state_to_db(session, ratchet_state_db, dr_state)
|
||||||
|
|
||||||
|
plaintext = plaintext_bytes.decode()
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Decrypted message: %s <- %s (msg #%s)",
|
||||||
|
recipient_id,
|
||||||
|
sender_id,
|
||||||
|
encrypted_msg.message_number
|
||||||
|
)
|
||||||
|
|
||||||
|
return plaintext
|
||||||
|
|
||||||
|
|
||||||
|
message_service = MessageService()
|
||||||
|
|
@ -0,0 +1,360 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Prekey management service for X3DH key bundles
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from datetime import (
|
||||||
|
UTC,
|
||||||
|
datetime,
|
||||||
|
timedelta,
|
||||||
|
)
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from sqlmodel import select
|
||||||
|
from sqlalchemy.exc import IntegrityError
|
||||||
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
DEFAULT_ONE_TIME_PREKEY_COUNT,
|
||||||
|
SIGNED_PREKEY_RETENTION_DAYS,
|
||||||
|
SIGNED_PREKEY_ROTATION_HOURS,
|
||||||
|
)
|
||||||
|
from app.core.encryption.x3dh_manager import (
|
||||||
|
PreKeyBundle,
|
||||||
|
x3dh_manager,
|
||||||
|
)
|
||||||
|
from app.core.exceptions import (
|
||||||
|
DatabaseError,
|
||||||
|
InvalidDataError,
|
||||||
|
UserNotFoundError,
|
||||||
|
)
|
||||||
|
from app.models.User import User
|
||||||
|
from app.models.IdentityKey import IdentityKey
|
||||||
|
from app.models.SignedPrekey import SignedPrekey
|
||||||
|
from app.models.OneTimePrekey import OneTimePrekey
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class PrekeyService:
|
||||||
|
"""
|
||||||
|
Service for managing X3DH prekey bundles and key rotation
|
||||||
|
"""
|
||||||
|
async def initialize_user_keys(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
user_id: UUID
|
||||||
|
) -> IdentityKey:
|
||||||
|
"""
|
||||||
|
Generates and stores initial identity key,
|
||||||
|
signed prekey, and one time prekeys for a user
|
||||||
|
"""
|
||||||
|
statement = select(User).where(User.id == user_id)
|
||||||
|
result = await session.execute(statement)
|
||||||
|
user = result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if not user:
|
||||||
|
logger.error("User not found: %s", user_id)
|
||||||
|
raise UserNotFoundError("User not found")
|
||||||
|
|
||||||
|
existing_ik_statement = select(IdentityKey).where(
|
||||||
|
IdentityKey.user_id == user_id
|
||||||
|
)
|
||||||
|
existing_ik_result = await session.execute(existing_ik_statement)
|
||||||
|
existing_ik = existing_ik_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if existing_ik:
|
||||||
|
logger.warning("Identity key already exists for user %s", user_id)
|
||||||
|
return existing_ik
|
||||||
|
|
||||||
|
ik_private_x25519, ik_public_x25519 = (
|
||||||
|
x3dh_manager.generate_identity_keypair_x25519()
|
||||||
|
)
|
||||||
|
ik_private_ed25519, ik_public_ed25519 = (
|
||||||
|
x3dh_manager.generate_identity_keypair_ed25519()
|
||||||
|
)
|
||||||
|
|
||||||
|
identity_key = IdentityKey(
|
||||||
|
user_id = user_id,
|
||||||
|
public_key = ik_public_x25519,
|
||||||
|
private_key = ik_private_x25519,
|
||||||
|
public_key_ed25519 = ik_public_ed25519,
|
||||||
|
private_key_ed25519 = ik_private_ed25519
|
||||||
|
)
|
||||||
|
|
||||||
|
session.add(identity_key)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
await session.refresh(identity_key)
|
||||||
|
logger.info("Created identity key for user %s", user_id)
|
||||||
|
except IntegrityError as e:
|
||||||
|
await session.rollback()
|
||||||
|
logger.error("Database error creating identity key: %s", e)
|
||||||
|
raise DatabaseError("Failed to create identity key") from e
|
||||||
|
|
||||||
|
await self.rotate_signed_prekey(session, user_id)
|
||||||
|
|
||||||
|
await self.replenish_one_time_prekeys(
|
||||||
|
session,
|
||||||
|
user_id,
|
||||||
|
DEFAULT_ONE_TIME_PREKEY_COUNT
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Initialized all keys for user %s: IK + SPK + %s OPKs",
|
||||||
|
user_id,
|
||||||
|
DEFAULT_ONE_TIME_PREKEY_COUNT
|
||||||
|
)
|
||||||
|
|
||||||
|
return identity_key
|
||||||
|
|
||||||
|
async def rotate_signed_prekey(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
user_id: UUID
|
||||||
|
) -> SignedPrekey:
|
||||||
|
"""
|
||||||
|
Generates new signed prekey and marks old ones inactive
|
||||||
|
"""
|
||||||
|
ik_statement = select(IdentityKey).where(IdentityKey.user_id == user_id)
|
||||||
|
ik_result = await session.execute(ik_statement)
|
||||||
|
identity_key = ik_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if not identity_key:
|
||||||
|
logger.error("Identity key not found for user %s", user_id)
|
||||||
|
raise InvalidDataError("User has no identity key")
|
||||||
|
|
||||||
|
old_spks_statement = select(SignedPrekey).where(
|
||||||
|
SignedPrekey.user_id == user_id,
|
||||||
|
SignedPrekey.is_active
|
||||||
|
)
|
||||||
|
old_spks_result = await session.execute(old_spks_statement)
|
||||||
|
old_spks = old_spks_result.scalars().all()
|
||||||
|
|
||||||
|
for old_spk in old_spks:
|
||||||
|
old_spk.is_active = False
|
||||||
|
logger.debug("Marked SPK %s as inactive", old_spk.key_id)
|
||||||
|
|
||||||
|
max_key_id_statement = select(SignedPrekey.key_id).where(
|
||||||
|
SignedPrekey.user_id == user_id
|
||||||
|
).order_by(SignedPrekey.key_id.desc()).limit(1)
|
||||||
|
max_key_id_result = await session.execute(max_key_id_statement)
|
||||||
|
max_key_id = max_key_id_result.scalar_one_or_none()
|
||||||
|
new_key_id = (max_key_id + 1) if max_key_id is not None else 1
|
||||||
|
|
||||||
|
spk_private, spk_public, spk_signature = (
|
||||||
|
x3dh_manager.generate_signed_prekey(
|
||||||
|
identity_key.private_key_ed25519
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
expires_at = datetime.now(UTC) + timedelta(
|
||||||
|
hours = SIGNED_PREKEY_ROTATION_HOURS
|
||||||
|
)
|
||||||
|
|
||||||
|
signed_prekey = SignedPrekey(
|
||||||
|
user_id = user_id,
|
||||||
|
key_id = new_key_id,
|
||||||
|
public_key = spk_public,
|
||||||
|
private_key = spk_private,
|
||||||
|
signature = spk_signature,
|
||||||
|
is_active = True,
|
||||||
|
expires_at = expires_at
|
||||||
|
)
|
||||||
|
|
||||||
|
session.add(signed_prekey)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
await session.refresh(signed_prekey)
|
||||||
|
logger.info(
|
||||||
|
"Rotated signed prekey for user %s: key_id=%s, expires=%s",
|
||||||
|
user_id,
|
||||||
|
new_key_id,
|
||||||
|
expires_at
|
||||||
|
)
|
||||||
|
except IntegrityError as e:
|
||||||
|
await session.rollback()
|
||||||
|
logger.error("Database error rotating signed prekey: %s", e)
|
||||||
|
raise DatabaseError("Failed to rotate signed prekey") from e
|
||||||
|
|
||||||
|
return signed_prekey
|
||||||
|
|
||||||
|
async def get_prekey_bundle(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
user_id: UUID
|
||||||
|
) -> PreKeyBundle:
|
||||||
|
"""
|
||||||
|
Retrieves prekey bundle for initiating X3DH with a user
|
||||||
|
"""
|
||||||
|
ik_statement = select(IdentityKey).where(IdentityKey.user_id == user_id)
|
||||||
|
ik_result = await session.execute(ik_statement)
|
||||||
|
identity_key = ik_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if not identity_key:
|
||||||
|
logger.error("Identity key not found for user %s", user_id)
|
||||||
|
raise InvalidDataError("User has no identity key")
|
||||||
|
|
||||||
|
spk_statement = select(SignedPrekey).where(
|
||||||
|
SignedPrekey.user_id == user_id,
|
||||||
|
SignedPrekey.is_active
|
||||||
|
).order_by(SignedPrekey.created_at.desc())
|
||||||
|
spk_result = await session.execute(spk_statement)
|
||||||
|
signed_prekey = spk_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
if not signed_prekey:
|
||||||
|
logger.warning(
|
||||||
|
"No active signed prekey for user %s, rotating",
|
||||||
|
user_id
|
||||||
|
)
|
||||||
|
signed_prekey = await self.rotate_signed_prekey(session, user_id)
|
||||||
|
|
||||||
|
opk_statement = select(OneTimePrekey).where(
|
||||||
|
OneTimePrekey.user_id == user_id,
|
||||||
|
not OneTimePrekey.is_used
|
||||||
|
).limit(1)
|
||||||
|
opk_result = await session.execute(opk_statement)
|
||||||
|
one_time_prekey = opk_result.scalar_one_or_none()
|
||||||
|
|
||||||
|
one_time_prekey_public = None
|
||||||
|
if one_time_prekey:
|
||||||
|
one_time_prekey.is_used = True
|
||||||
|
one_time_prekey_public = one_time_prekey.public_key
|
||||||
|
logger.debug(
|
||||||
|
"Consumed one time prekey %s for user %s",
|
||||||
|
one_time_prekey.key_id,
|
||||||
|
user_id
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
except IntegrityError as e:
|
||||||
|
await session.rollback()
|
||||||
|
logger.error("Database error consuming OPK: %s", e)
|
||||||
|
raise DatabaseError("Failed to consume one-time prekey") from e
|
||||||
|
|
||||||
|
bundle = PreKeyBundle(
|
||||||
|
identity_key = identity_key.public_key,
|
||||||
|
signed_prekey = signed_prekey.public_key,
|
||||||
|
signed_prekey_signature = signed_prekey.signature,
|
||||||
|
one_time_prekey = one_time_prekey_public
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Retrieved prekey bundle for user %s: IK + SPK + %s",
|
||||||
|
user_id,
|
||||||
|
'OPK' if one_time_prekey_public else 'no OPK'
|
||||||
|
)
|
||||||
|
|
||||||
|
return bundle
|
||||||
|
|
||||||
|
async def replenish_one_time_prekeys(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
user_id: UUID,
|
||||||
|
count: int = DEFAULT_ONE_TIME_PREKEY_COUNT
|
||||||
|
) -> int:
|
||||||
|
"""
|
||||||
|
Generates new batch of one time prekeys
|
||||||
|
"""
|
||||||
|
max_key_id_statement = select(OneTimePrekey.key_id).where(
|
||||||
|
OneTimePrekey.user_id == user_id
|
||||||
|
).order_by(OneTimePrekey.key_id.desc()).limit(1)
|
||||||
|
max_key_id_result = await session.execute(max_key_id_statement)
|
||||||
|
max_key_id = max_key_id_result.scalar_one_or_none()
|
||||||
|
next_key_id = (max_key_id + 1) if max_key_id is not None else 1
|
||||||
|
|
||||||
|
one_time_prekeys = []
|
||||||
|
for i in range(count):
|
||||||
|
opk_private, opk_public = x3dh_manager.generate_one_time_prekey()
|
||||||
|
|
||||||
|
one_time_prekey = OneTimePrekey(
|
||||||
|
user_id = user_id,
|
||||||
|
key_id = next_key_id + i,
|
||||||
|
public_key = opk_public,
|
||||||
|
private_key = opk_private,
|
||||||
|
is_used = False
|
||||||
|
)
|
||||||
|
one_time_prekeys.append(one_time_prekey)
|
||||||
|
|
||||||
|
for opk in one_time_prekeys:
|
||||||
|
session.add(opk)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
logger.info(
|
||||||
|
"Generated %s one-time prekeys for user %s",
|
||||||
|
count,
|
||||||
|
user_id
|
||||||
|
)
|
||||||
|
except IntegrityError as e:
|
||||||
|
await session.rollback()
|
||||||
|
logger.error("Database error generating OPKs: %s", e)
|
||||||
|
raise DatabaseError("Failed to generate one-time prekeys") from e
|
||||||
|
|
||||||
|
return count
|
||||||
|
|
||||||
|
async def get_unused_opk_count(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
user_id: UUID
|
||||||
|
) -> int:
|
||||||
|
"""
|
||||||
|
Returns count of unused one time prekeys for a user
|
||||||
|
"""
|
||||||
|
count_statement = select(OneTimePrekey).where(
|
||||||
|
OneTimePrekey.user_id == user_id,
|
||||||
|
not OneTimePrekey.is_used
|
||||||
|
)
|
||||||
|
result = await session.execute(count_statement)
|
||||||
|
unused_opks = result.scalars().all()
|
||||||
|
|
||||||
|
count = len(unused_opks)
|
||||||
|
logger.debug("User %s has %s unused OPKs", user_id, count)
|
||||||
|
return count
|
||||||
|
|
||||||
|
async def cleanup_old_signed_prekeys(
|
||||||
|
self,
|
||||||
|
session: AsyncSession,
|
||||||
|
user_id: UUID
|
||||||
|
) -> int:
|
||||||
|
"""
|
||||||
|
Deletes inactive signed prekeys older than retention period
|
||||||
|
"""
|
||||||
|
cutoff_date = datetime.now(UTC) - timedelta(
|
||||||
|
days = SIGNED_PREKEY_RETENTION_DAYS
|
||||||
|
)
|
||||||
|
|
||||||
|
old_spks_statement = select(SignedPrekey).where(
|
||||||
|
SignedPrekey.user_id == user_id,
|
||||||
|
not SignedPrekey.is_active,
|
||||||
|
SignedPrekey.created_at < cutoff_date
|
||||||
|
)
|
||||||
|
old_spks_result = await session.execute(old_spks_statement)
|
||||||
|
old_spks = old_spks_result.scalars().all()
|
||||||
|
|
||||||
|
deleted_count = len(old_spks)
|
||||||
|
|
||||||
|
for spk in old_spks:
|
||||||
|
await session.delete(spk)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
logger.info(
|
||||||
|
"Deleted %s old signed prekeys for user %s",
|
||||||
|
deleted_count,
|
||||||
|
user_id
|
||||||
|
)
|
||||||
|
except IntegrityError as e:
|
||||||
|
await session.rollback()
|
||||||
|
logger.error("Database error deleting old SPKs: %s", e)
|
||||||
|
raise DatabaseError("Failed to delete old signed prekeys") from e
|
||||||
|
|
||||||
|
return deleted_count
|
||||||
|
|
||||||
|
|
||||||
|
prekey_service = PrekeyService()
|
||||||
|
|
@ -0,0 +1,164 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Presence service for managing user online/offline status
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from app.core.enums import PresenceStatus
|
||||||
|
from app.core.exceptions import DatabaseError
|
||||||
|
from app.core.surreal_manager import surreal_db
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class PresenceService:
|
||||||
|
"""
|
||||||
|
Service for managing user presence status in real time
|
||||||
|
"""
|
||||||
|
async def set_user_online(self, user_id: UUID) -> None:
|
||||||
|
"""
|
||||||
|
Mark user as online and update last seen timestamp
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
await surreal_db.update_presence(
|
||||||
|
user_id = str(user_id),
|
||||||
|
status = PresenceStatus.ONLINE.value,
|
||||||
|
last_seen = datetime.now(UTC).isoformat()
|
||||||
|
)
|
||||||
|
logger.info("User %s is now online", user_id)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to set user %s online: %s", user_id, e)
|
||||||
|
raise DatabaseError(f"Failed to update presence: {str(e)}") from e
|
||||||
|
|
||||||
|
async def set_user_offline(self, user_id: UUID) -> None:
|
||||||
|
"""
|
||||||
|
Mark user as offline and update last seen timestamp
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
await surreal_db.update_presence(
|
||||||
|
user_id = str(user_id),
|
||||||
|
status = PresenceStatus.OFFLINE.value,
|
||||||
|
last_seen = datetime.now(UTC).isoformat()
|
||||||
|
)
|
||||||
|
logger.info("User %s is now offline", user_id)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to set user %s offline: %s", user_id, e)
|
||||||
|
raise DatabaseError(f"Failed to update presence: {str(e)}") from e
|
||||||
|
|
||||||
|
async def set_user_away(self, user_id: UUID) -> None:
|
||||||
|
"""
|
||||||
|
Mark user as away due to inactivity
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
await surreal_db.update_presence(
|
||||||
|
user_id = str(user_id),
|
||||||
|
status = PresenceStatus.AWAY.value,
|
||||||
|
last_seen = datetime.now(UTC).isoformat()
|
||||||
|
)
|
||||||
|
logger.debug("User %s is now away", user_id)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to set user %s away: %s", user_id, e)
|
||||||
|
raise DatabaseError(f"Failed to update presence: {str(e)}") from e
|
||||||
|
|
||||||
|
async def update_last_seen(self, user_id: UUID) -> None:
|
||||||
|
"""
|
||||||
|
Update user last seen timestamp without changing status
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
last_seen = datetime.now(UTC).isoformat()
|
||||||
|
await surreal_db.db.merge(
|
||||||
|
f"presence:{user_id}",
|
||||||
|
{
|
||||||
|
"last_seen": last_seen,
|
||||||
|
"updated_at": "time::now()"
|
||||||
|
}
|
||||||
|
)
|
||||||
|
logger.debug("Updated last seen for user %s", user_id)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to update last seen for %s: %s", user_id, e)
|
||||||
|
|
||||||
|
async def get_user_presence(self, user_id: UUID) -> dict:
|
||||||
|
"""
|
||||||
|
Get current presence status for a user
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
await surreal_db.ensure_connected()
|
||||||
|
result = await surreal_db.db.select(f"presence:{user_id}")
|
||||||
|
|
||||||
|
if not result:
|
||||||
|
return {
|
||||||
|
"user_id": str(user_id),
|
||||||
|
"status": PresenceStatus.OFFLINE.value,
|
||||||
|
"last_seen": datetime.now(UTC).isoformat()
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
"user_id": result.get("user_id",
|
||||||
|
str(user_id)),
|
||||||
|
"status": result.get("status",
|
||||||
|
PresenceStatus.OFFLINE.value),
|
||||||
|
"last_seen":
|
||||||
|
result.get("last_seen",
|
||||||
|
datetime.now(UTC).isoformat())
|
||||||
|
}
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to get presence for user %s: %s", user_id, e)
|
||||||
|
return {
|
||||||
|
"user_id": str(user_id),
|
||||||
|
"status": PresenceStatus.OFFLINE.value,
|
||||||
|
"last_seen": datetime.now(UTC).isoformat()
|
||||||
|
}
|
||||||
|
|
||||||
|
async def get_room_online_users(self, room_id: str) -> list[dict]:
|
||||||
|
"""
|
||||||
|
Get all online users in a specific room
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
presence_list = await surreal_db.get_room_presence(room_id)
|
||||||
|
|
||||||
|
return [
|
||||||
|
{
|
||||||
|
"user_id": p.user_id,
|
||||||
|
"status": p.status,
|
||||||
|
"last_seen": p.last_seen.isoformat()
|
||||||
|
} for p in presence_list
|
||||||
|
]
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to get online users for room %s: %s", room_id, e)
|
||||||
|
return []
|
||||||
|
|
||||||
|
async def bulk_update_presence(
|
||||||
|
self,
|
||||||
|
user_ids: list[UUID],
|
||||||
|
status: PresenceStatus
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Update presence status for multiple users at once
|
||||||
|
"""
|
||||||
|
for user_id in user_ids:
|
||||||
|
try:
|
||||||
|
await surreal_db.update_presence(
|
||||||
|
user_id = str(user_id),
|
||||||
|
status = status.value,
|
||||||
|
last_seen = datetime.now(UTC).isoformat()
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
"Failed to bulk update presence for %s: %s",
|
||||||
|
user_id,
|
||||||
|
e
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Bulk updated presence for %s users to %s",
|
||||||
|
len(user_ids),
|
||||||
|
status.value
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
presence_service = PresenceService()
|
||||||
|
|
@ -0,0 +1,293 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
WebSocket service for handling real time message routing and processing
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from typing import Any
|
||||||
|
from uuid import UUID
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
|
||||||
|
from fastapi import WebSocket
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
WS_MESSAGE_TYPE_ENCRYPTED,
|
||||||
|
WS_MESSAGE_TYPE_PRESENCE,
|
||||||
|
WS_MESSAGE_TYPE_RECEIPT,
|
||||||
|
WS_MESSAGE_TYPE_TYPING,
|
||||||
|
)
|
||||||
|
from app.core.enums import PresenceStatus
|
||||||
|
from app.core.websocket_manager import connection_manager
|
||||||
|
from app.schemas.websocket import (
|
||||||
|
EncryptedMessageWS,
|
||||||
|
ReadReceiptWS,
|
||||||
|
TypingIndicatorWS,
|
||||||
|
)
|
||||||
|
from app.models.Base import async_session_maker
|
||||||
|
from app.services.message_service import message_service
|
||||||
|
from app.services.presence_service import presence_service
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class WebSocketService:
|
||||||
|
"""
|
||||||
|
Service for processing WebSocket
|
||||||
|
messages and routing to appropriate handlers
|
||||||
|
"""
|
||||||
|
async def route_message(
|
||||||
|
self,
|
||||||
|
websocket: WebSocket,
|
||||||
|
user_id: UUID,
|
||||||
|
message: dict[str,
|
||||||
|
Any]
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Route incoming WebSocket message
|
||||||
|
to appropriate handler based on type
|
||||||
|
"""
|
||||||
|
message_type = message.get("type")
|
||||||
|
|
||||||
|
if not message_type:
|
||||||
|
await websocket.send_json(
|
||||||
|
{
|
||||||
|
"type": "error",
|
||||||
|
"error_code": "missing_type",
|
||||||
|
"error_message": "Message type is required"
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
if message_type == WS_MESSAGE_TYPE_ENCRYPTED:
|
||||||
|
await self.handle_encrypted_message(user_id, message)
|
||||||
|
elif message_type == WS_MESSAGE_TYPE_TYPING:
|
||||||
|
await self.handle_typing_indicator(user_id, message)
|
||||||
|
elif message_type == WS_MESSAGE_TYPE_PRESENCE:
|
||||||
|
await self.handle_presence_update(user_id, message)
|
||||||
|
elif message_type == WS_MESSAGE_TYPE_RECEIPT:
|
||||||
|
await self.handle_read_receipt(user_id, message)
|
||||||
|
elif message_type == "heartbeat":
|
||||||
|
await self.handle_heartbeat(user_id)
|
||||||
|
else:
|
||||||
|
logger.warning(
|
||||||
|
"Unknown message type from %s: %s",
|
||||||
|
user_id,
|
||||||
|
message_type
|
||||||
|
)
|
||||||
|
await websocket.send_json(
|
||||||
|
{
|
||||||
|
"type": "error",
|
||||||
|
"error_code": "unknown_type",
|
||||||
|
"error_message": f"Unknown message type: {message_type}"
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
async def handle_encrypted_message(
|
||||||
|
self,
|
||||||
|
user_id: UUID,
|
||||||
|
message: dict[str,
|
||||||
|
Any]
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Process encrypted message from client and forward to recipient
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
recipient_id = UUID(message.get("recipient_id"))
|
||||||
|
plaintext = message.get("plaintext")
|
||||||
|
|
||||||
|
if not plaintext:
|
||||||
|
logger.error("Missing plaintext in message from %s", user_id)
|
||||||
|
return
|
||||||
|
|
||||||
|
async with async_session_maker() as session:
|
||||||
|
result = await message_service.send_encrypted_message(
|
||||||
|
session,
|
||||||
|
user_id,
|
||||||
|
recipient_id,
|
||||||
|
plaintext
|
||||||
|
)
|
||||||
|
|
||||||
|
ws_message = EncryptedMessageWS(
|
||||||
|
message_id = result.id if hasattr(result,
|
||||||
|
'id') else "unknown",
|
||||||
|
sender_id = str(user_id),
|
||||||
|
recipient_id = str(recipient_id),
|
||||||
|
ciphertext = message.get("ciphertext",
|
||||||
|
""),
|
||||||
|
nonce = message.get("nonce",
|
||||||
|
""),
|
||||||
|
header = message.get("header",
|
||||||
|
""),
|
||||||
|
sender_username = message.get("sender_username",
|
||||||
|
"")
|
||||||
|
)
|
||||||
|
|
||||||
|
await connection_manager.send_message(
|
||||||
|
recipient_id,
|
||||||
|
ws_message.model_dump(mode = "json")
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Encrypted message forwarded: %s -> %s",
|
||||||
|
user_id,
|
||||||
|
recipient_id
|
||||||
|
)
|
||||||
|
|
||||||
|
except ValueError as e:
|
||||||
|
logger.error(
|
||||||
|
"Invalid UUID in encrypted message from %s: %s",
|
||||||
|
user_id,
|
||||||
|
e
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
"Failed to handle encrypted message from %s: %s",
|
||||||
|
user_id,
|
||||||
|
e
|
||||||
|
)
|
||||||
|
|
||||||
|
async def handle_typing_indicator(
|
||||||
|
self,
|
||||||
|
user_id: UUID,
|
||||||
|
message: dict[str,
|
||||||
|
Any]
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Process typing indicator and broadcast to room
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
room_id = message.get("room_id")
|
||||||
|
is_typing = message.get("is_typing", False)
|
||||||
|
|
||||||
|
if not room_id:
|
||||||
|
logger.error(
|
||||||
|
"Missing room_id in typing indicator from %s",
|
||||||
|
user_id
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
typing_msg = TypingIndicatorWS(
|
||||||
|
user_id = str(user_id),
|
||||||
|
room_id = room_id,
|
||||||
|
is_typing = is_typing
|
||||||
|
)
|
||||||
|
|
||||||
|
await connection_manager.broadcast_to_room(
|
||||||
|
room_id,
|
||||||
|
typing_msg.model_dump(mode = "json")
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug(
|
||||||
|
"Typing indicator broadcast: %s in %s = %s",
|
||||||
|
user_id,
|
||||||
|
room_id,
|
||||||
|
is_typing
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
"Failed to handle typing indicator from %s: %s",
|
||||||
|
user_id,
|
||||||
|
e
|
||||||
|
)
|
||||||
|
|
||||||
|
async def handle_presence_update(
|
||||||
|
self,
|
||||||
|
user_id: UUID,
|
||||||
|
message: dict[str,
|
||||||
|
Any]
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Process presence status update from client
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
status = message.get("status")
|
||||||
|
|
||||||
|
if not status:
|
||||||
|
logger.error("Missing status in presence update from %s", user_id)
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
presence_status = PresenceStatus(status)
|
||||||
|
except ValueError:
|
||||||
|
logger.warning(
|
||||||
|
"Invalid presence status from %s: %s",
|
||||||
|
user_id,
|
||||||
|
status
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
if presence_status == PresenceStatus.ONLINE:
|
||||||
|
await presence_service.set_user_online(user_id)
|
||||||
|
elif presence_status == PresenceStatus.AWAY:
|
||||||
|
await presence_service.set_user_away(user_id)
|
||||||
|
elif presence_status == PresenceStatus.OFFLINE:
|
||||||
|
await presence_service.set_user_offline(user_id)
|
||||||
|
|
||||||
|
logger.debug(
|
||||||
|
"Presence updated: %s -> %s",
|
||||||
|
user_id,
|
||||||
|
presence_status.value
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
"Failed to handle presence update from %s: %s",
|
||||||
|
user_id,
|
||||||
|
e
|
||||||
|
)
|
||||||
|
|
||||||
|
async def handle_read_receipt(
|
||||||
|
self,
|
||||||
|
user_id: UUID,
|
||||||
|
message: dict[str,
|
||||||
|
Any]
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Process read receipt and notify message sender
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
message_id = message.get("message_id")
|
||||||
|
sender_id_str = message.get("sender_id")
|
||||||
|
|
||||||
|
if not message_id or not sender_id_str:
|
||||||
|
logger.error(
|
||||||
|
"Missing message_id or sender_id in receipt from %s",
|
||||||
|
user_id
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
sender_id = UUID(sender_id_str)
|
||||||
|
|
||||||
|
receipt_msg = ReadReceiptWS(
|
||||||
|
message_id = message_id,
|
||||||
|
user_id = str(user_id),
|
||||||
|
read_at = datetime.now(UTC)
|
||||||
|
)
|
||||||
|
|
||||||
|
await connection_manager.send_message(
|
||||||
|
sender_id,
|
||||||
|
receipt_msg.model_dump(mode = "json")
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug(
|
||||||
|
"Read receipt sent: message %s read by %s",
|
||||||
|
message_id,
|
||||||
|
user_id
|
||||||
|
)
|
||||||
|
|
||||||
|
except ValueError as e:
|
||||||
|
logger.error("Invalid UUID in read receipt from %s: %s", user_id, e)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("Failed to handle read receipt from %s: %s", user_id, e)
|
||||||
|
|
||||||
|
async def handle_heartbeat(self, user_id: UUID) -> None:
|
||||||
|
"""
|
||||||
|
Process heartbeat message and update user last seen
|
||||||
|
"""
|
||||||
|
logger.debug("Heartbeat received from user %s", user_id)
|
||||||
|
await presence_service.update_last_seen(user_id)
|
||||||
|
|
||||||
|
|
||||||
|
websocket_service = WebSocketService()
|
||||||
|
|
@ -0,0 +1,249 @@
|
||||||
|
[project]
|
||||||
|
name = "encrypted-p2p-chat"
|
||||||
|
version = "1.0.0"
|
||||||
|
description = "End-to-end encrypted P2P chat with Triple Ratchet and WebAuthn"
|
||||||
|
requires-python = ">=3.13"
|
||||||
|
authors = [
|
||||||
|
{name = "Carter", email = "carter@certgames.com"}
|
||||||
|
]
|
||||||
|
|
||||||
|
dependencies = [
|
||||||
|
"fastapi>=0.121.0",
|
||||||
|
"uvicorn[standard]>=0.38.0",
|
||||||
|
"websockets>=15.0.1",
|
||||||
|
"redis[hiredis]>=7.1.0",
|
||||||
|
"sqlalchemy>=2.0.44",
|
||||||
|
"sqlmodel>=0.0.27",
|
||||||
|
"alembic>=1.17.2",
|
||||||
|
"asyncpg>=0.30.0",
|
||||||
|
"pydantic>=2.12.4",
|
||||||
|
"pydantic-settings>=2.12.0",
|
||||||
|
"webauthn>=2.7.0",
|
||||||
|
"fido2>=2.0.0",
|
||||||
|
"cryptography>=46.0.3",
|
||||||
|
"pynacl>=1.6.1",
|
||||||
|
"passlib>=1.7.4",
|
||||||
|
"python-multipart>=0.0.20",
|
||||||
|
"httpx>=0.28.1",
|
||||||
|
"orjson>=3.11.4",
|
||||||
|
"surrealdb>=1.0.6",
|
||||||
|
"liboqs-python>=0.14.1"
|
||||||
|
]
|
||||||
|
[project.optional-dependencies]
|
||||||
|
dev = [
|
||||||
|
"pytest>=8.0.0",
|
||||||
|
"pytest-asyncio>=0.23.0",
|
||||||
|
"aiosqlite>=0.20.0",
|
||||||
|
"ruff>=0.8.0",
|
||||||
|
"mypy>=1.9.0",
|
||||||
|
"pre-commit>=3.0.0",
|
||||||
|
"types-redis>=4.6.0",
|
||||||
|
"types-passlib>=1.7.0",
|
||||||
|
]
|
||||||
|
|
||||||
|
[build-system]
|
||||||
|
requires = ["setuptools>=80.9.0", "wheel>=0.45.1"]
|
||||||
|
build-backend = "setuptools.build_meta"
|
||||||
|
|
||||||
|
[tool.setuptools.packages.find]
|
||||||
|
include = ["app"]
|
||||||
|
|
||||||
|
[tool.ruff]
|
||||||
|
target-version = "py313"
|
||||||
|
line-length = 95
|
||||||
|
indent-width = 4
|
||||||
|
exclude = [
|
||||||
|
".bzr",
|
||||||
|
".direnv",
|
||||||
|
".eggs",
|
||||||
|
".git",
|
||||||
|
".git-rewrite",
|
||||||
|
".hg",
|
||||||
|
".ipynb_checkpoints",
|
||||||
|
".mypy_cache",
|
||||||
|
".nox",
|
||||||
|
".pyenv",
|
||||||
|
".pytest_cache",
|
||||||
|
".pytype",
|
||||||
|
".ruff_cache",
|
||||||
|
".svn",
|
||||||
|
".tox",
|
||||||
|
".venv",
|
||||||
|
"__pycache__",
|
||||||
|
"_build",
|
||||||
|
"build",
|
||||||
|
"dist",
|
||||||
|
"site-packages",
|
||||||
|
"venv",
|
||||||
|
]
|
||||||
|
|
||||||
|
[tool.ruff.format]
|
||||||
|
line-ending = "auto"
|
||||||
|
skip-magic-trailing-comma = false
|
||||||
|
|
||||||
|
[tool.ruff.lint]
|
||||||
|
select = [
|
||||||
|
"E1", # Indentation
|
||||||
|
"E4", # Imports
|
||||||
|
"E7", # Statement
|
||||||
|
"F", # Pyflakes (all F rules)
|
||||||
|
"W292", # No newline at end of file
|
||||||
|
"W605", # Invalid escape sequence
|
||||||
|
"B", # Bugbear
|
||||||
|
"C4", # Comprehensions
|
||||||
|
"UP", # Pyupgrade
|
||||||
|
"ARG", # Unused arguments
|
||||||
|
"SIM", # Simplify
|
||||||
|
"I", # isort rules
|
||||||
|
"F401", # Unused imports
|
||||||
|
"F811", # Redefined imports
|
||||||
|
"F821", # Undefined name
|
||||||
|
]
|
||||||
|
|
||||||
|
ignore = [
|
||||||
|
"E501", # Line length (handled by formatter)
|
||||||
|
"W291", # Trailing whitespace
|
||||||
|
"W293", # Blank line contains whitespace
|
||||||
|
"I001", # Import sorting
|
||||||
|
"RUF001", # Ambiguous unicode
|
||||||
|
"RUF002", # Docstring with ambiguous unicode
|
||||||
|
"B008", # FastAPI Depends() in defaults is standard pattern
|
||||||
|
"ARG001", # Unused function args (lifespan protocol)
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
[tool.ruff.lint.per-file-ignores]
|
||||||
|
"__init__.py" = ["F401"]
|
||||||
|
|
||||||
|
[tool.mypy]
|
||||||
|
python_version = "3.13"
|
||||||
|
strict = true
|
||||||
|
warn_return_any = true
|
||||||
|
warn_unused_configs = true
|
||||||
|
disallow_untyped_defs = true
|
||||||
|
disallow_any_generics = true
|
||||||
|
check_untyped_defs = true
|
||||||
|
no_implicit_optional = true
|
||||||
|
warn_redundant_casts = true
|
||||||
|
warn_unused_ignores = true
|
||||||
|
warn_no_return = true
|
||||||
|
follow_imports = "normal"
|
||||||
|
ignore_missing_imports = true
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "app.models.*"
|
||||||
|
disable_error_code = ["call-arg", "misc"]
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "app.core.surreal_manager"
|
||||||
|
disable_error_code = ["union-attr", "no-any-return"]
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "app.core.redis_manager"
|
||||||
|
disable_error_code = ["type-arg", "attr-defined"]
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "app.core.exception_handlers"
|
||||||
|
disable_error_code = ["arg-type"]
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "app.services.presence_service"
|
||||||
|
disable_error_code = ["union-attr", "type-arg"]
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "app.services.auth_service"
|
||||||
|
disable_error_code = ["arg-type", "no-any-return"]
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "app.services.prekey_service"
|
||||||
|
disable_error_code = ["attr-defined", "no-any-return"]
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "app.services.message_service"
|
||||||
|
disable_error_code = ["no-any-return"]
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = "app.core.websocket_manager"
|
||||||
|
disable_error_code = ["attr-defined", "type-arg"]
|
||||||
|
|
||||||
|
|
||||||
|
[tool.pylint.main]
|
||||||
|
py-version = "3.13"
|
||||||
|
jobs = 4
|
||||||
|
load-plugins = [
|
||||||
|
"pylint_mongoengine",
|
||||||
|
"pylint_pydantic",
|
||||||
|
"pylint_per_file_ignores",
|
||||||
|
"pylint_flask",
|
||||||
|
"pylint_celery"
|
||||||
|
]
|
||||||
|
persistent = true
|
||||||
|
suggestion-mode = true
|
||||||
|
ignore = [
|
||||||
|
"venv",
|
||||||
|
".venv",
|
||||||
|
"__pycache__",
|
||||||
|
"build",
|
||||||
|
"dist",
|
||||||
|
".git",
|
||||||
|
".pytest_cache",
|
||||||
|
".mypy_cache",
|
||||||
|
".ruff_cache",
|
||||||
|
]
|
||||||
|
ignore-paths = [
|
||||||
|
"^venv/.*",
|
||||||
|
"^.venv/.*",
|
||||||
|
"^build/.*",
|
||||||
|
"^dist/.*",
|
||||||
|
]
|
||||||
|
[tool.pylint.type-check]
|
||||||
|
generated-members = [
|
||||||
|
"objects",
|
||||||
|
"id",
|
||||||
|
"get_or_create",
|
||||||
|
"DoesNotExist",
|
||||||
|
"MultipleObjectsReturned",
|
||||||
|
"objects.get_or_create"
|
||||||
|
]
|
||||||
|
|
||||||
|
[tool.pylint.messages_control]
|
||||||
|
disable = [
|
||||||
|
"C0111", # missing-docstring
|
||||||
|
"C0103", # invalid-name
|
||||||
|
"R0903", # too-few-public-methods
|
||||||
|
"W0511", # fixme
|
||||||
|
"W0622", # redefined-builtin
|
||||||
|
"W0612", # unused-variable (handled by ruff)
|
||||||
|
"W0613", # unused-argument (handled by ruff)
|
||||||
|
"C0301", # Line too long
|
||||||
|
"C0302", # Too many lines
|
||||||
|
"C0411", # Wrong import order
|
||||||
|
"C0305", # Trailing newlines
|
||||||
|
"C0303", # Trailing whitespace
|
||||||
|
"C0304", # Final newline missing
|
||||||
|
"R0801", # Similar lines (want exact duplicates only)
|
||||||
|
"E0401", # Unable to import - packages not in pylint env
|
||||||
|
"C0412", # Import grouping - don't care about grouping imports
|
||||||
|
"W0718", # Broad exception catching - intentional for service layer
|
||||||
|
"E0611", # No name in module - false positive for dynamic imports
|
||||||
|
"E1101", # No member - false positive for alembic context
|
||||||
|
]
|
||||||
|
|
||||||
|
[tool.pylint.design]
|
||||||
|
max-args = 7
|
||||||
|
max-attributes = 10
|
||||||
|
max-locals = 30
|
||||||
|
max-positional-arguments = 7
|
||||||
|
|
||||||
|
[tool.pylint."messages control"]
|
||||||
|
per-file-ignores = [
|
||||||
|
"alembic/env.py:W0611", # Unused imports needed for SQLModel metadata
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
[tool.pytest.ini_options]
|
||||||
|
asyncio_mode = "auto"
|
||||||
|
testpaths = ["tests"]
|
||||||
|
python_files = ["test_*.py"]
|
||||||
|
python_classes = ["Test*"]
|
||||||
|
python_functions = ["test_*"]
|
||||||
|
|
@ -0,0 +1,4 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
encrypted-p2p-chat pyest suite
|
||||||
|
"""
|
||||||
|
|
@ -0,0 +1,209 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Pytest configuration and fixtures for all tests
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from uuid import uuid4
|
||||||
|
from typing import Any
|
||||||
|
from collections.abc import AsyncGenerator
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import pytest_asyncio
|
||||||
|
from sqlmodel import SQLModel
|
||||||
|
from sqlalchemy.ext.asyncio import (
|
||||||
|
AsyncSession,
|
||||||
|
create_async_engine,
|
||||||
|
)
|
||||||
|
from sqlalchemy.orm import sessionmaker
|
||||||
|
from webauthn.helpers import bytes_to_base64url
|
||||||
|
|
||||||
|
from app.models.User import User
|
||||||
|
from app.models.IdentityKey import IdentityKey
|
||||||
|
from app.models.SignedPrekey import SignedPrekey
|
||||||
|
from app.models.OneTimePrekey import OneTimePrekey
|
||||||
|
from app.core.encryption.x3dh_manager import x3dh_manager
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope = "session")
|
||||||
|
def event_loop():
|
||||||
|
"""
|
||||||
|
Create event loop for async tests
|
||||||
|
"""
|
||||||
|
loop = asyncio.get_event_loop_policy().new_event_loop()
|
||||||
|
yield loop
|
||||||
|
loop.close()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture(scope = "function")
|
||||||
|
async def db_session() -> AsyncGenerator[AsyncSession]:
|
||||||
|
"""
|
||||||
|
Create in-memory SQLite database for testing
|
||||||
|
"""
|
||||||
|
engine = create_async_engine(
|
||||||
|
"sqlite+aiosqlite:///:memory:",
|
||||||
|
echo = False,
|
||||||
|
future = True,
|
||||||
|
)
|
||||||
|
|
||||||
|
async with engine.begin() as conn:
|
||||||
|
await conn.run_sync(SQLModel.metadata.create_all)
|
||||||
|
|
||||||
|
async_session = sessionmaker(
|
||||||
|
bind = engine,
|
||||||
|
class_ = AsyncSession,
|
||||||
|
expire_on_commit = False,
|
||||||
|
)
|
||||||
|
|
||||||
|
async with async_session() as session:
|
||||||
|
yield session
|
||||||
|
await session.rollback()
|
||||||
|
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
async def test_user(db_session: AsyncSession) -> User:
|
||||||
|
"""
|
||||||
|
Create test user
|
||||||
|
"""
|
||||||
|
user = User(
|
||||||
|
username = "testuser",
|
||||||
|
display_name = "Test User",
|
||||||
|
is_active = True,
|
||||||
|
is_verified = True,
|
||||||
|
)
|
||||||
|
db_session.add(user)
|
||||||
|
await db_session.commit()
|
||||||
|
await db_session.refresh(user)
|
||||||
|
return user
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
async def test_user_2(db_session: AsyncSession) -> User:
|
||||||
|
"""
|
||||||
|
Create second test user for conversations
|
||||||
|
"""
|
||||||
|
user = User(
|
||||||
|
username = "testuser2",
|
||||||
|
display_name = "Test User 2",
|
||||||
|
is_active = True,
|
||||||
|
is_verified = True,
|
||||||
|
)
|
||||||
|
db_session.add(user)
|
||||||
|
await db_session.commit()
|
||||||
|
await db_session.refresh(user)
|
||||||
|
return user
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
async def test_identity_key(
|
||||||
|
db_session: AsyncSession,
|
||||||
|
test_user: User
|
||||||
|
) -> IdentityKey:
|
||||||
|
"""
|
||||||
|
Create identity key for test user
|
||||||
|
"""
|
||||||
|
ik_private_x25519, ik_public_x25519 = (
|
||||||
|
x3dh_manager.generate_identity_keypair_x25519()
|
||||||
|
)
|
||||||
|
ik_private_ed25519, ik_public_ed25519 = (
|
||||||
|
x3dh_manager.generate_identity_keypair_ed25519()
|
||||||
|
)
|
||||||
|
|
||||||
|
identity_key = IdentityKey(
|
||||||
|
user_id = test_user.id,
|
||||||
|
public_key = ik_public_x25519,
|
||||||
|
private_key = ik_private_x25519,
|
||||||
|
public_key_ed25519 = ik_public_ed25519,
|
||||||
|
private_key_ed25519 = ik_private_ed25519,
|
||||||
|
)
|
||||||
|
|
||||||
|
db_session.add(identity_key)
|
||||||
|
await db_session.commit()
|
||||||
|
await db_session.refresh(identity_key)
|
||||||
|
return identity_key
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
async def test_signed_prekey(
|
||||||
|
db_session: AsyncSession,
|
||||||
|
test_user: User,
|
||||||
|
test_identity_key: IdentityKey
|
||||||
|
) -> SignedPrekey:
|
||||||
|
"""
|
||||||
|
Create signed prekey for test user
|
||||||
|
"""
|
||||||
|
spk_private, spk_public, spk_signature = x3dh_manager.generate_signed_prekey(
|
||||||
|
test_identity_key.private_key_ed25519
|
||||||
|
)
|
||||||
|
|
||||||
|
signed_prekey = SignedPrekey(
|
||||||
|
user_id = test_user.id,
|
||||||
|
key_id = 1,
|
||||||
|
public_key = spk_public,
|
||||||
|
private_key = spk_private,
|
||||||
|
signature = spk_signature,
|
||||||
|
is_active = True,
|
||||||
|
)
|
||||||
|
|
||||||
|
db_session.add(signed_prekey)
|
||||||
|
await db_session.commit()
|
||||||
|
await db_session.refresh(signed_prekey)
|
||||||
|
return signed_prekey
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture
|
||||||
|
async def test_one_time_prekey(
|
||||||
|
db_session: AsyncSession,
|
||||||
|
test_user: User
|
||||||
|
) -> OneTimePrekey:
|
||||||
|
"""
|
||||||
|
Create one-time prekey for test user
|
||||||
|
"""
|
||||||
|
opk_private, opk_public = x3dh_manager.generate_one_time_prekey()
|
||||||
|
|
||||||
|
one_time_prekey = OneTimePrekey(
|
||||||
|
user_id = test_user.id,
|
||||||
|
key_id = 1,
|
||||||
|
public_key = opk_public,
|
||||||
|
private_key = opk_private,
|
||||||
|
is_used = False,
|
||||||
|
)
|
||||||
|
|
||||||
|
db_session.add(one_time_prekey)
|
||||||
|
await db_session.commit()
|
||||||
|
await db_session.refresh(one_time_prekey)
|
||||||
|
return one_time_prekey
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_webauthn_credential() -> dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Mock WebAuthn credential response
|
||||||
|
"""
|
||||||
|
return {
|
||||||
|
"id": bytes_to_base64url(uuid4().bytes),
|
||||||
|
"rawId": bytes_to_base64url(uuid4().bytes),
|
||||||
|
"type": "public-key",
|
||||||
|
"response": {
|
||||||
|
"clientDataJSON": bytes_to_base64url(b'{"type":"webauthn.create"}'),
|
||||||
|
"attestationObject": bytes_to_base64url(b"mock_attestation"),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def sample_plaintext() -> str:
|
||||||
|
"""
|
||||||
|
Sample message for encryption tests
|
||||||
|
"""
|
||||||
|
return "Hello, this is a test message!"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def sample_associated_data() -> bytes:
|
||||||
|
"""
|
||||||
|
Sample associated data for AEAD
|
||||||
|
"""
|
||||||
|
return b"test_sender:test_recipient"
|
||||||
|
|
@ -0,0 +1,97 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Tests for authentication service
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
|
from app.models.User import User
|
||||||
|
from app.services.auth_service import auth_service
|
||||||
|
from app.core.exceptions import UserExistsError
|
||||||
|
|
||||||
|
|
||||||
|
class TestAuthService:
|
||||||
|
"""
|
||||||
|
Test authentication service basics
|
||||||
|
"""
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_create_user(self, db_session: AsyncSession):
|
||||||
|
"""
|
||||||
|
Test creating a new user
|
||||||
|
"""
|
||||||
|
user = await auth_service.create_user(
|
||||||
|
session = db_session,
|
||||||
|
username = "newuser",
|
||||||
|
display_name = "New User"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert user.id is not None
|
||||||
|
assert user.username == "newuser"
|
||||||
|
assert user.display_name == "New User"
|
||||||
|
assert user.is_active is True
|
||||||
|
assert user.is_verified is False
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_create_duplicate_user_fails(
|
||||||
|
self,
|
||||||
|
db_session: AsyncSession,
|
||||||
|
test_user: User
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Test cannot create user with duplicate username
|
||||||
|
"""
|
||||||
|
with pytest.raises(UserExistsError, match = "already exists"):
|
||||||
|
await auth_service.create_user(
|
||||||
|
session = db_session,
|
||||||
|
username = test_user.username,
|
||||||
|
display_name = "Duplicate"
|
||||||
|
)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_user_by_username(
|
||||||
|
self,
|
||||||
|
db_session: AsyncSession,
|
||||||
|
test_user: User
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Test retrieving user by username
|
||||||
|
"""
|
||||||
|
user = await auth_service.get_user_by_username(
|
||||||
|
session = db_session,
|
||||||
|
username = test_user.username
|
||||||
|
)
|
||||||
|
|
||||||
|
assert user is not None
|
||||||
|
assert user.id == test_user.id
|
||||||
|
assert user.username == test_user.username
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_nonexistent_user(self, db_session: AsyncSession):
|
||||||
|
"""
|
||||||
|
Test getting user that doesn't exist returns None
|
||||||
|
"""
|
||||||
|
user = await auth_service.get_user_by_username(
|
||||||
|
session = db_session,
|
||||||
|
username = "nonexistent"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert user is None
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_user_by_id(
|
||||||
|
self,
|
||||||
|
db_session: AsyncSession,
|
||||||
|
test_user: User
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Test retrieving user by ID
|
||||||
|
"""
|
||||||
|
user = await auth_service.get_user_by_id(
|
||||||
|
session = db_session,
|
||||||
|
user_id = test_user.id
|
||||||
|
)
|
||||||
|
|
||||||
|
assert user is not None
|
||||||
|
assert user.id == test_user.id
|
||||||
|
assert user.username == test_user.username
|
||||||
|
|
@ -0,0 +1,165 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Tests for Double Ratchet encryption core
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from cryptography.hazmat.primitives import serialization
|
||||||
|
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
|
||||||
|
|
||||||
|
from app.core.encryption.double_ratchet import double_ratchet
|
||||||
|
|
||||||
|
|
||||||
|
class TestDoubleRatchet:
|
||||||
|
"""
|
||||||
|
Test Double Ratchet encryption/decryption
|
||||||
|
"""
|
||||||
|
def test_encrypt_decrypt_basic(
|
||||||
|
self,
|
||||||
|
sample_plaintext: str,
|
||||||
|
sample_associated_data: bytes
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Test basic encrypt/decrypt cycle works
|
||||||
|
"""
|
||||||
|
shared_key = b"0" * 32
|
||||||
|
|
||||||
|
bob_dh_private = X25519PrivateKey.generate()
|
||||||
|
bob_dh_public_bytes = bob_dh_private.public_key().public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
|
||||||
|
sender_state = double_ratchet.initialize_sender(
|
||||||
|
shared_key = shared_key,
|
||||||
|
peer_public_key = bob_dh_public_bytes
|
||||||
|
)
|
||||||
|
|
||||||
|
receiver_state = double_ratchet.initialize_receiver(
|
||||||
|
shared_key = shared_key,
|
||||||
|
own_private_key = bob_dh_private
|
||||||
|
)
|
||||||
|
|
||||||
|
encrypted = double_ratchet.encrypt_message(
|
||||||
|
sender_state,
|
||||||
|
sample_plaintext.encode(),
|
||||||
|
sample_associated_data
|
||||||
|
)
|
||||||
|
|
||||||
|
decrypted = double_ratchet.decrypt_message(
|
||||||
|
receiver_state,
|
||||||
|
encrypted,
|
||||||
|
sample_associated_data
|
||||||
|
)
|
||||||
|
|
||||||
|
assert decrypted.decode() == sample_plaintext
|
||||||
|
|
||||||
|
def test_multiple_messages(self, sample_associated_data: bytes):
|
||||||
|
"""
|
||||||
|
Test multiple messages maintain state correctly
|
||||||
|
"""
|
||||||
|
shared_key = b"0" * 32
|
||||||
|
|
||||||
|
bob_dh_private = X25519PrivateKey.generate()
|
||||||
|
bob_dh_public_bytes = bob_dh_private.public_key().public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
|
||||||
|
sender_state = double_ratchet.initialize_sender(
|
||||||
|
shared_key = shared_key,
|
||||||
|
peer_public_key = bob_dh_public_bytes
|
||||||
|
)
|
||||||
|
|
||||||
|
receiver_state = double_ratchet.initialize_receiver(
|
||||||
|
shared_key = shared_key,
|
||||||
|
own_private_key = bob_dh_private
|
||||||
|
)
|
||||||
|
|
||||||
|
messages = [b"Message 1", b"Message 2", b"Message 3"]
|
||||||
|
encrypted_messages = []
|
||||||
|
|
||||||
|
for msg in messages:
|
||||||
|
encrypted = double_ratchet.encrypt_message(
|
||||||
|
sender_state,
|
||||||
|
msg,
|
||||||
|
sample_associated_data
|
||||||
|
)
|
||||||
|
encrypted_messages.append(encrypted)
|
||||||
|
|
||||||
|
for i, encrypted in enumerate(encrypted_messages):
|
||||||
|
decrypted = double_ratchet.decrypt_message(
|
||||||
|
receiver_state,
|
||||||
|
encrypted,
|
||||||
|
sample_associated_data
|
||||||
|
)
|
||||||
|
assert decrypted == messages[i]
|
||||||
|
|
||||||
|
def test_message_numbers_increment(self, sample_associated_data: bytes):
|
||||||
|
"""
|
||||||
|
Test message numbers increment correctly
|
||||||
|
"""
|
||||||
|
shared_key = b"0" * 32
|
||||||
|
peer_public_key = b"1" * 32
|
||||||
|
|
||||||
|
sender_state = double_ratchet.initialize_sender(
|
||||||
|
shared_key = shared_key,
|
||||||
|
peer_public_key = peer_public_key
|
||||||
|
)
|
||||||
|
|
||||||
|
assert sender_state.sending_message_number == 0
|
||||||
|
|
||||||
|
double_ratchet.encrypt_message(
|
||||||
|
sender_state,
|
||||||
|
b"Message 1",
|
||||||
|
sample_associated_data
|
||||||
|
)
|
||||||
|
|
||||||
|
assert sender_state.sending_message_number == 1
|
||||||
|
|
||||||
|
double_ratchet.encrypt_message(
|
||||||
|
sender_state,
|
||||||
|
b"Message 2",
|
||||||
|
sample_associated_data
|
||||||
|
)
|
||||||
|
|
||||||
|
assert sender_state.sending_message_number == 2
|
||||||
|
|
||||||
|
def test_tampered_message_fails(self, sample_associated_data: bytes):
|
||||||
|
"""
|
||||||
|
Test tampered messages fail to decrypt
|
||||||
|
"""
|
||||||
|
shared_key = b"0" * 32
|
||||||
|
|
||||||
|
bob_dh_private = X25519PrivateKey.generate()
|
||||||
|
bob_dh_public_bytes = bob_dh_private.public_key().public_bytes(
|
||||||
|
encoding = serialization.Encoding.Raw,
|
||||||
|
format = serialization.PublicFormat.Raw
|
||||||
|
)
|
||||||
|
|
||||||
|
sender_state = double_ratchet.initialize_sender(
|
||||||
|
shared_key = shared_key,
|
||||||
|
peer_public_key = bob_dh_public_bytes
|
||||||
|
)
|
||||||
|
|
||||||
|
receiver_state = double_ratchet.initialize_receiver(
|
||||||
|
shared_key = shared_key,
|
||||||
|
own_private_key = bob_dh_private
|
||||||
|
)
|
||||||
|
|
||||||
|
encrypted = double_ratchet.encrypt_message(
|
||||||
|
sender_state,
|
||||||
|
b"Original message",
|
||||||
|
sample_associated_data
|
||||||
|
)
|
||||||
|
|
||||||
|
tampered_ciphertext = bytearray(encrypted.ciphertext)
|
||||||
|
tampered_ciphertext[0] ^= 0xFF
|
||||||
|
encrypted.ciphertext = bytes(tampered_ciphertext)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match = "tampered or corrupted"):
|
||||||
|
double_ratchet.decrypt_message(
|
||||||
|
receiver_state,
|
||||||
|
encrypted,
|
||||||
|
sample_associated_data
|
||||||
|
)
|
||||||
|
|
@ -0,0 +1,135 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Tests for message service (end to end encryption flow)
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
|
from app.models.User import User
|
||||||
|
from app.models.IdentityKey import IdentityKey
|
||||||
|
from app.models.SignedPrekey import SignedPrekey
|
||||||
|
from app.core.exceptions import InvalidDataError
|
||||||
|
from app.models.OneTimePrekey import OneTimePrekey
|
||||||
|
from app.services.message_service import message_service
|
||||||
|
from app.core.encryption.x3dh_manager import x3dh_manager
|
||||||
|
|
||||||
|
|
||||||
|
class TestMessageService:
|
||||||
|
"""
|
||||||
|
Test message encryption/decryption service
|
||||||
|
"""
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_initialize_conversation(
|
||||||
|
self,
|
||||||
|
db_session: AsyncSession,
|
||||||
|
test_user: User,
|
||||||
|
test_user_2: User,
|
||||||
|
test_identity_key: IdentityKey,
|
||||||
|
test_signed_prekey: SignedPrekey,
|
||||||
|
test_one_time_prekey: OneTimePrekey
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Test initializing encrypted conversation between two users
|
||||||
|
"""
|
||||||
|
sender_ik_private, sender_ik_public = (
|
||||||
|
x3dh_manager.generate_identity_keypair_x25519()
|
||||||
|
)
|
||||||
|
sender_ik_private_ed, sender_ik_public_ed = (
|
||||||
|
x3dh_manager.generate_identity_keypair_ed25519()
|
||||||
|
)
|
||||||
|
|
||||||
|
sender_identity_key = IdentityKey(
|
||||||
|
user_id = test_user_2.id,
|
||||||
|
public_key = sender_ik_public,
|
||||||
|
private_key = sender_ik_private,
|
||||||
|
public_key_ed25519 = sender_ik_public_ed,
|
||||||
|
private_key_ed25519 = sender_ik_private_ed,
|
||||||
|
)
|
||||||
|
db_session.add(sender_identity_key)
|
||||||
|
await db_session.commit()
|
||||||
|
|
||||||
|
ratchet_state = await message_service.initialize_conversation(
|
||||||
|
session = db_session,
|
||||||
|
sender_id = test_user_2.id,
|
||||||
|
recipient_id = test_user.id
|
||||||
|
)
|
||||||
|
|
||||||
|
assert ratchet_state.user_id == test_user_2.id
|
||||||
|
assert ratchet_state.peer_user_id == test_user.id
|
||||||
|
assert ratchet_state.sending_message_number == 0
|
||||||
|
assert ratchet_state.receiving_message_number == 0
|
||||||
|
assert ratchet_state.dh_public_key is not None
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_cannot_initialize_with_self(
|
||||||
|
self,
|
||||||
|
db_session: AsyncSession,
|
||||||
|
test_user: User
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Test cannot start conversation with yourself
|
||||||
|
"""
|
||||||
|
with pytest.raises(InvalidDataError,
|
||||||
|
match = "Cannot start conversation with yourself"):
|
||||||
|
await message_service.initialize_conversation(
|
||||||
|
session = db_session,
|
||||||
|
sender_id = test_user.id,
|
||||||
|
recipient_id = test_user.id
|
||||||
|
)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_ratchet_state_persistence(
|
||||||
|
self,
|
||||||
|
db_session: AsyncSession,
|
||||||
|
test_user: User,
|
||||||
|
test_user_2: User,
|
||||||
|
test_identity_key: IdentityKey,
|
||||||
|
test_signed_prekey: SignedPrekey,
|
||||||
|
test_one_time_prekey: OneTimePrekey
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Test ratchet state loads and saves correctly
|
||||||
|
"""
|
||||||
|
sender_ik_private, sender_ik_public = (
|
||||||
|
x3dh_manager.generate_identity_keypair_x25519()
|
||||||
|
)
|
||||||
|
sender_ik_private_ed, sender_ik_public_ed = (
|
||||||
|
x3dh_manager.generate_identity_keypair_ed25519()
|
||||||
|
)
|
||||||
|
|
||||||
|
sender_identity_key = IdentityKey(
|
||||||
|
user_id = test_user_2.id,
|
||||||
|
public_key = sender_ik_public,
|
||||||
|
private_key = sender_ik_private,
|
||||||
|
public_key_ed25519 = sender_ik_public_ed,
|
||||||
|
private_key_ed25519 = sender_ik_private_ed,
|
||||||
|
)
|
||||||
|
db_session.add(sender_identity_key)
|
||||||
|
await db_session.commit()
|
||||||
|
|
||||||
|
ratchet_state = await message_service.initialize_conversation(
|
||||||
|
session = db_session,
|
||||||
|
sender_id = test_user_2.id,
|
||||||
|
recipient_id = test_user.id
|
||||||
|
)
|
||||||
|
|
||||||
|
initial_msg_num = ratchet_state.sending_message_number
|
||||||
|
|
||||||
|
dr_state = await message_service._load_ratchet_state_from_db(
|
||||||
|
ratchet_state
|
||||||
|
)
|
||||||
|
|
||||||
|
assert dr_state.sending_message_number == initial_msg_num
|
||||||
|
assert dr_state.dh_private_key is not None
|
||||||
|
|
||||||
|
dr_state.sending_message_number += 1
|
||||||
|
|
||||||
|
await message_service._save_ratchet_state_to_db(
|
||||||
|
db_session,
|
||||||
|
ratchet_state,
|
||||||
|
dr_state
|
||||||
|
)
|
||||||
|
|
||||||
|
await db_session.refresh(ratchet_state)
|
||||||
|
assert ratchet_state.sending_message_number == initial_msg_num + 1
|
||||||
|
|
@ -0,0 +1,130 @@
|
||||||
|
"""
|
||||||
|
ⒸAngelaMos | 2025
|
||||||
|
Tests for X3DH key exchange protocol
|
||||||
|
"""
|
||||||
|
|
||||||
|
from webauthn.helpers import base64url_to_bytes
|
||||||
|
|
||||||
|
from app.core.encryption.x3dh_manager import x3dh_manager, PreKeyBundle
|
||||||
|
|
||||||
|
|
||||||
|
class TestX3DH:
|
||||||
|
"""
|
||||||
|
Test X3DH key exchange
|
||||||
|
"""
|
||||||
|
def test_key_generation(self):
|
||||||
|
"""
|
||||||
|
Test all key generation functions work
|
||||||
|
"""
|
||||||
|
ik_private, ik_public = x3dh_manager.generate_identity_keypair_x25519()
|
||||||
|
assert len(base64url_to_bytes(ik_private)) == 32
|
||||||
|
assert len(base64url_to_bytes(ik_public)) == 32
|
||||||
|
|
||||||
|
ik_private_ed, ik_public_ed = x3dh_manager.generate_identity_keypair_ed25519()
|
||||||
|
assert len(base64url_to_bytes(ik_private_ed)) == 32
|
||||||
|
assert len(base64url_to_bytes(ik_public_ed)) == 32
|
||||||
|
|
||||||
|
spk_private, spk_public, signature = x3dh_manager.generate_signed_prekey(
|
||||||
|
ik_private_ed
|
||||||
|
)
|
||||||
|
assert len(base64url_to_bytes(spk_private)) == 32
|
||||||
|
assert len(base64url_to_bytes(spk_public)) == 32
|
||||||
|
assert len(base64url_to_bytes(signature)) == 64
|
||||||
|
|
||||||
|
opk_private, opk_public = x3dh_manager.generate_one_time_prekey()
|
||||||
|
assert len(base64url_to_bytes(opk_private)) == 32
|
||||||
|
assert len(base64url_to_bytes(opk_public)) == 32
|
||||||
|
|
||||||
|
def test_x3dh_handshake_with_opk(self):
|
||||||
|
"""
|
||||||
|
Test full X3DH handshake with one-time prekey
|
||||||
|
"""
|
||||||
|
alice_ik_private, alice_ik_public = (
|
||||||
|
x3dh_manager.generate_identity_keypair_x25519()
|
||||||
|
)
|
||||||
|
alice_ik_private_ed, alice_ik_public_ed = (
|
||||||
|
x3dh_manager.generate_identity_keypair_ed25519()
|
||||||
|
)
|
||||||
|
|
||||||
|
bob_ik_private, bob_ik_public = (
|
||||||
|
x3dh_manager.generate_identity_keypair_x25519()
|
||||||
|
)
|
||||||
|
bob_ik_private_ed, bob_ik_public_ed = (
|
||||||
|
x3dh_manager.generate_identity_keypair_ed25519()
|
||||||
|
)
|
||||||
|
|
||||||
|
bob_spk_private, bob_spk_public, bob_spk_sig = (
|
||||||
|
x3dh_manager.generate_signed_prekey(bob_ik_private_ed)
|
||||||
|
)
|
||||||
|
|
||||||
|
bob_opk_private, bob_opk_public = x3dh_manager.generate_one_time_prekey()
|
||||||
|
|
||||||
|
bob_bundle = PreKeyBundle(
|
||||||
|
identity_key = bob_ik_public,
|
||||||
|
signed_prekey = bob_spk_public,
|
||||||
|
signed_prekey_signature = bob_spk_sig,
|
||||||
|
one_time_prekey = bob_opk_public
|
||||||
|
)
|
||||||
|
|
||||||
|
alice_result = x3dh_manager.perform_x3dh_sender(
|
||||||
|
alice_identity_private_x25519 = alice_ik_private,
|
||||||
|
bob_bundle = bob_bundle,
|
||||||
|
bob_identity_public_ed25519 = bob_ik_public_ed
|
||||||
|
)
|
||||||
|
|
||||||
|
bob_result = x3dh_manager.perform_x3dh_receiver(
|
||||||
|
bob_identity_private_x25519 = bob_ik_private,
|
||||||
|
bob_signed_prekey_private = bob_spk_private,
|
||||||
|
bob_one_time_prekey_private = bob_opk_private,
|
||||||
|
alice_ephemeral_public = alice_result.ephemeral_public_key,
|
||||||
|
alice_identity_public_x25519 = alice_ik_public
|
||||||
|
)
|
||||||
|
|
||||||
|
assert alice_result.shared_key == bob_result.shared_key
|
||||||
|
assert len(alice_result.shared_key) == 32
|
||||||
|
|
||||||
|
def test_x3dh_handshake_without_opk(self):
|
||||||
|
"""
|
||||||
|
Test X3DH handshake without one-time prekey
|
||||||
|
"""
|
||||||
|
alice_ik_private, alice_ik_public = (
|
||||||
|
x3dh_manager.generate_identity_keypair_x25519()
|
||||||
|
)
|
||||||
|
alice_ik_private_ed, alice_ik_public_ed = (
|
||||||
|
x3dh_manager.generate_identity_keypair_ed25519()
|
||||||
|
)
|
||||||
|
|
||||||
|
bob_ik_private, bob_ik_public = (
|
||||||
|
x3dh_manager.generate_identity_keypair_x25519()
|
||||||
|
)
|
||||||
|
bob_ik_private_ed, bob_ik_public_ed = (
|
||||||
|
x3dh_manager.generate_identity_keypair_ed25519()
|
||||||
|
)
|
||||||
|
|
||||||
|
bob_spk_private, bob_spk_public, bob_spk_sig = (
|
||||||
|
x3dh_manager.generate_signed_prekey(bob_ik_private_ed)
|
||||||
|
)
|
||||||
|
|
||||||
|
bob_bundle = PreKeyBundle(
|
||||||
|
identity_key = bob_ik_public,
|
||||||
|
signed_prekey = bob_spk_public,
|
||||||
|
signed_prekey_signature = bob_spk_sig,
|
||||||
|
one_time_prekey = None
|
||||||
|
)
|
||||||
|
|
||||||
|
alice_result = x3dh_manager.perform_x3dh_sender(
|
||||||
|
alice_identity_private_x25519 = alice_ik_private,
|
||||||
|
bob_bundle = bob_bundle,
|
||||||
|
bob_identity_public_ed25519 = bob_ik_public_ed
|
||||||
|
)
|
||||||
|
|
||||||
|
bob_result = x3dh_manager.perform_x3dh_receiver(
|
||||||
|
bob_identity_private_x25519 = bob_ik_private,
|
||||||
|
bob_signed_prekey_private = bob_spk_private,
|
||||||
|
bob_one_time_prekey_private = None,
|
||||||
|
alice_ephemeral_public = alice_result.ephemeral_public_key,
|
||||||
|
alice_identity_public_x25519 = alice_ik_public
|
||||||
|
)
|
||||||
|
|
||||||
|
assert alice_result.shared_key == bob_result.shared_key
|
||||||
|
assert len(alice_result.shared_key) == 32
|
||||||
Loading…
Reference in New Issue