Files
Kecalek_python/zaloha/db.py
2026-03-11 16:54:14 +01:00

1294 lines
43 KiB
Python

"""MySQL database layer for the encrypted chat server."""
import os
import uuid
import mysql.connector
from dotenv import load_dotenv
from crypto_utils import (
generate_identity_keypair,
serialize_ed25519_public,
generate_signed_prekey,
serialize_x25519_public,
generate_one_time_prekeys,
)
load_dotenv()
# Sentinel device_id for self-encrypted copies and legacy (pre-multi-device) rows
SELF_DEVICE_ID = "00000000-0000-0000-0000-000000000000"
def get_connection():
"""Create a new MySQL connection from environment variables."""
return mysql.connector.connect(
host=os.getenv("MYSQL_HOST", "localhost"),
port=int(os.getenv("MYSQL_PORT", "3306")),
user=os.getenv("MYSQL_USER", "root"),
password=os.getenv("MYSQL_PASSWORD", ""),
database=os.getenv("MYSQL_DATABASE", "encrypted_chat"),
)
def generate_uuid() -> str:
return str(uuid.uuid4())
# --- Devices ---
def create_device(user_id: str, device_name: str | None = None) -> str:
"""Create a new device for a user. Returns device_id."""
conn = get_connection()
try:
cursor = conn.cursor()
device_id = generate_uuid()
cursor.execute(
"INSERT INTO devices (id, user_id, device_name) VALUES (%s, %s, %s)",
(device_id, user_id, device_name),
)
conn.commit()
return device_id
finally:
conn.close()
def get_user_devices(user_id: str) -> list[dict]:
"""Get all devices for a user."""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT id, user_id, device_name, created_at, last_seen_at "
"FROM devices WHERE user_id = %s ORDER BY created_at",
(user_id,),
)
return cursor.fetchall()
finally:
conn.close()
def get_device(device_id: str) -> dict | None:
"""Get a single device by ID."""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT id, user_id, device_name, created_at, last_seen_at "
"FROM devices WHERE id = %s",
(device_id,),
)
return cursor.fetchone()
finally:
conn.close()
def update_device_last_seen(device_id: str):
"""Update last_seen_at timestamp for a device."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"UPDATE devices SET last_seen_at = NOW() WHERE id = %s",
(device_id,),
)
conn.commit()
finally:
conn.close()
def delete_device(device_id: str):
"""Delete a device. CASCADE removes its prekeys."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute("DELETE FROM devices WHERE id = %s", (device_id,))
# Also clean up prekeys explicitly for device_id column
cursor.execute("DELETE FROM signed_prekeys WHERE device_id = %s", (device_id,))
cursor.execute("DELETE FROM one_time_prekeys WHERE device_id = %s", (device_id,))
conn.commit()
finally:
conn.close()
# --- Users ---
def create_user(username: str, email: str, rsa_public_key_pem: str, identity_key: bytes) -> str:
"""Register a new user. Returns user ID."""
conn = get_connection()
try:
cursor = conn.cursor()
user_id = generate_uuid()
cursor.execute(
"INSERT INTO users (id, username, email, rsa_public_key, identity_key) "
"VALUES (%s, %s, %s, %s, %s)",
(user_id, username, email, rsa_public_key_pem, identity_key),
)
conn.commit()
return user_id
finally:
conn.close()
def get_user_by_email(email: str) -> dict | None:
"""Get user by email."""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT id, username, rsa_public_key, email, identity_key FROM users WHERE email = %s",
(email,),
)
return cursor.fetchone()
finally:
conn.close()
def get_user_by_id(user_id: str) -> dict | None:
"""Get user by ID."""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT id, username, rsa_public_key, email, identity_key FROM users WHERE id = %s",
(user_id,),
)
return cursor.fetchone()
finally:
conn.close()
def get_user_contacts(user_id: str) -> list[str]:
"""Get all user IDs that share at least one conversation with the given user."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"SELECT DISTINCT cm2.user_id "
"FROM conversation_members cm1 "
"JOIN conversation_members cm2 ON cm1.conversation_id = cm2.conversation_id "
"WHERE cm1.user_id = %s AND cm2.user_id != %s",
(user_id, user_id),
)
return [row[0] for row in cursor.fetchall()]
finally:
conn.close()
def update_user_rsa_key(user_id: str, rsa_public_key_pem: str):
"""Update user's RSA public key (for login)."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute("UPDATE users SET rsa_public_key = %s WHERE id = %s", (rsa_public_key_pem, user_id))
conn.commit()
finally:
conn.close()
# --- Pre-keys ---
def store_signed_prekey(user_id: str, spk_id: str, public_key: bytes, signature: bytes,
device_id: str | None = None):
"""Store (or replace) a signed pre-key for a user's device."""
conn = get_connection()
try:
cursor = conn.cursor()
# Remove old SPKs for this user+device
if device_id:
cursor.execute("DELETE FROM signed_prekeys WHERE user_id = %s AND device_id = %s",
(user_id, device_id))
else:
cursor.execute("DELETE FROM signed_prekeys WHERE user_id = %s AND device_id IS NULL",
(user_id,))
cursor.execute(
"INSERT INTO signed_prekeys (id, user_id, device_id, public_key, signature) "
"VALUES (%s, %s, %s, %s, %s)",
(spk_id, user_id, device_id, public_key, signature),
)
conn.commit()
finally:
conn.close()
def get_signed_prekey(user_id: str, device_id: str | None = None) -> dict | None:
"""Get the current signed pre-key for a user (optionally per device)."""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
if device_id:
cursor.execute(
"SELECT id, public_key, signature, device_id, created_at FROM signed_prekeys "
"WHERE user_id = %s AND device_id = %s "
"ORDER BY created_at DESC LIMIT 1",
(user_id, device_id),
)
else:
cursor.execute(
"SELECT id, public_key, signature, device_id, created_at FROM signed_prekeys "
"WHERE user_id = %s ORDER BY created_at DESC LIMIT 1",
(user_id,),
)
return cursor.fetchone()
finally:
conn.close()
def store_one_time_prekeys(user_id: str, prekeys: list[dict], device_id: str | None = None):
"""Store a batch of one-time pre-keys. Each dict has {id, public_key (bytes)}."""
conn = get_connection()
try:
cursor = conn.cursor()
for pk in prekeys:
cursor.execute(
"INSERT INTO one_time_prekeys (id, user_id, device_id, public_key) "
"VALUES (%s, %s, %s, %s)",
(pk["id"], user_id, device_id, pk["public_key"]),
)
conn.commit()
finally:
conn.close()
def consume_one_time_prekey(user_id: str, device_id: str | None = None) -> dict | None:
"""Atomically consume one OTP: SELECT FOR UPDATE + DELETE.
Returns {id, public_key} or None."""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
conn.start_transaction()
if device_id:
cursor.execute(
"SELECT id, public_key FROM one_time_prekeys "
"WHERE user_id = %s AND device_id = %s LIMIT 1 FOR UPDATE",
(user_id, device_id),
)
else:
cursor.execute(
"SELECT id, public_key FROM one_time_prekeys "
"WHERE user_id = %s LIMIT 1 FOR UPDATE",
(user_id,),
)
row = cursor.fetchone()
if row:
cursor.execute("DELETE FROM one_time_prekeys WHERE id = %s", (row["id"],))
conn.commit()
return row
except Exception:
conn.rollback()
raise
finally:
conn.close()
def count_one_time_prekeys(user_id: str, device_id: str | None = None) -> int:
"""Count remaining OTPs for a user (optionally per device)."""
conn = get_connection()
try:
cursor = conn.cursor()
if device_id:
cursor.execute(
"SELECT COUNT(*) FROM one_time_prekeys WHERE user_id = %s AND device_id = %s",
(user_id, device_id),
)
else:
cursor.execute("SELECT COUNT(*) FROM one_time_prekeys WHERE user_id = %s", (user_id,))
return cursor.fetchone()[0]
finally:
conn.close()
def get_key_bundle(user_id: str) -> dict | None:
"""Get complete key bundle for X3DH (single device — legacy compat).
Returns {identity_key, signed_prekey_id, signed_prekey, spk_signature,
one_time_prekey_id, one_time_prekey} or None.
OTP is consumed atomically.
"""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
# Get user identity key
cursor.execute("SELECT identity_key FROM users WHERE id = %s", (user_id,))
user = cursor.fetchone()
if not user:
return None
# Get signed prekey
cursor.execute(
"SELECT id, public_key, signature, device_id FROM signed_prekeys WHERE user_id = %s "
"ORDER BY created_at DESC LIMIT 1",
(user_id,),
)
spk = cursor.fetchone()
if not spk:
return None
# Consume one OTP (may be None) — use transaction for atomicity (H12 fix)
conn.start_transaction()
cursor.execute(
"SELECT id, public_key FROM one_time_prekeys WHERE user_id = %s LIMIT 1 FOR UPDATE",
(user_id,),
)
opk = cursor.fetchone()
if opk:
cursor.execute("DELETE FROM one_time_prekeys WHERE id = %s", (opk["id"],))
conn.commit()
result = {
"identity_key": user["identity_key"],
"signed_prekey_id": spk["id"],
"signed_prekey": spk["public_key"],
"spk_signature": spk["signature"],
}
if opk:
result["one_time_prekey_id"] = opk["id"]
result["one_time_prekey"] = opk["public_key"]
return result
except Exception:
try:
conn.rollback()
except Exception:
pass
raise
finally:
conn.close()
def get_key_bundles_for_user(user_id: str) -> dict | None:
"""Get key bundles for ALL devices of a user. Returns
{identity_key, device_bundles: [{device_id, signed_prekey_id, signed_prekey_pub,
spk_signature, opk_id, opk_pub}]} or None.
Consumes one OPK per device atomically.
"""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
# Get user identity key
cursor.execute("SELECT identity_key FROM users WHERE id = %s", (user_id,))
user = cursor.fetchone()
if not user:
return None
# Get all signed prekeys (one per device, most recent)
cursor.execute(
"SELECT id, public_key, signature, device_id FROM signed_prekeys "
"WHERE user_id = %s ORDER BY created_at DESC",
(user_id,),
)
all_spks = cursor.fetchall()
if not all_spks:
return None
# De-duplicate: keep only the most recent SPK per device_id
seen_devices = set()
spks_by_device = []
for spk in all_spks:
dev = spk.get("device_id") or "__legacy__"
if dev not in seen_devices:
seen_devices.add(dev)
spks_by_device.append(spk)
device_bundles = []
# Commit the implicit transaction from the read-only queries above
# so we can start an explicit transaction for atomic OPK consumption.
conn.commit()
conn.start_transaction()
for spk in spks_by_device:
dev_id = spk.get("device_id")
# Consume one OPK for this device
if dev_id:
cursor.execute(
"SELECT id, public_key FROM one_time_prekeys "
"WHERE user_id = %s AND device_id = %s LIMIT 1 FOR UPDATE",
(user_id, dev_id),
)
else:
cursor.execute(
"SELECT id, public_key FROM one_time_prekeys "
"WHERE user_id = %s AND device_id IS NULL LIMIT 1 FOR UPDATE",
(user_id,),
)
opk = cursor.fetchone()
if opk:
cursor.execute("DELETE FROM one_time_prekeys WHERE id = %s", (opk["id"],))
bundle = {
"device_id": dev_id,
"signed_prekey_id": spk["id"],
"signed_prekey_pub": spk["public_key"],
"spk_signature": spk["signature"],
}
if opk:
bundle["opk_id"] = opk["id"]
bundle["opk_pub"] = opk["public_key"]
device_bundles.append(bundle)
conn.commit()
return {
"identity_key": user["identity_key"],
"device_bundles": device_bundles,
}
except Exception:
try:
conn.rollback()
except Exception:
pass
raise
finally:
conn.close()
# --- Conversations ---
def create_conversation(member_user_ids: list[str], joined_at=None, name=None, created_by=None) -> str:
conn = get_connection()
try:
cursor = conn.cursor()
conv_id = generate_uuid()
cursor.execute("INSERT INTO conversations (id, name, created_by) VALUES (%s, %s, %s)",
(conv_id, name, created_by))
for uid in member_user_ids:
cursor.execute(
"INSERT INTO conversation_members (conversation_id, user_id, joined_at) VALUES (%s, %s, %s)",
(conv_id, uid, joined_at),
)
conn.commit()
return conv_id
finally:
conn.close()
def add_conversation_member(conversation_id: str, user_id: str, joined_at=None):
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"INSERT IGNORE INTO conversation_members (conversation_id, user_id, joined_at) VALUES (%s, %s, %s)",
(conversation_id, user_id, joined_at),
)
conn.commit()
finally:
conn.close()
def remove_conversation_member(conversation_id: str, user_id: str):
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"DELETE FROM conversation_members WHERE conversation_id = %s AND user_id = %s",
(conversation_id, user_id),
)
conn.commit()
finally:
conn.close()
def count_conversation_members(conversation_id: str) -> int:
"""Count members in a conversation."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"SELECT COUNT(*) FROM conversation_members WHERE conversation_id = %s",
(conversation_id,),
)
return cursor.fetchone()[0]
finally:
conn.close()
def get_conversation_file_ids(conversation_id: str) -> list[str]:
"""Get all file IDs (images + files) uploaded to a conversation."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"SELECT file_id FROM image_uploads WHERE conversation_id = %s",
(conversation_id,),
)
return [row[0] for row in cursor.fetchall()]
finally:
conn.close()
def delete_conversation(conversation_id: str):
"""Delete a conversation entirely. CASCADE cleans up members, messages, etc."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute("DELETE FROM conversations WHERE id = %s", (conversation_id,))
conn.commit()
finally:
conn.close()
def get_conversation_members(conversation_id: str) -> list[dict]:
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT u.id, u.username, u.email, u.identity_key FROM conversation_members cm "
"JOIN users u ON cm.user_id = u.id "
"WHERE cm.conversation_id = %s",
(conversation_id,),
)
return cursor.fetchall()
finally:
conn.close()
def find_direct_conversation(user_id_a: str, user_id_b: str) -> str | None:
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"SELECT cm1.conversation_id FROM conversation_members cm1 "
"JOIN conversation_members cm2 ON cm1.conversation_id = cm2.conversation_id "
"WHERE cm1.user_id = %s AND cm2.user_id = %s "
"AND (SELECT COUNT(*) FROM conversation_members cm3 "
" WHERE cm3.conversation_id = cm1.conversation_id) = 2 "
"LIMIT 1",
(user_id_a, user_id_b),
)
row = cursor.fetchone()
return row[0] if row else None
finally:
conn.close()
def update_conversation_creator(conversation_id: str, new_creator_id: str):
"""Transfer group creator role to another member."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"UPDATE conversations SET created_by = %s WHERE id = %s",
(new_creator_id, conversation_id),
)
conn.commit()
finally:
conn.close()
def get_conversation(conversation_id: str) -> dict | None:
"""Get conversation by ID."""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT id, created_at, name, created_by, avatar_file FROM conversations WHERE id = %s",
(conversation_id,),
)
return cursor.fetchone()
finally:
conn.close()
def update_conversation_avatar(conversation_id: str, avatar_file: str):
"""Set avatar file for a conversation."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"UPDATE conversations SET avatar_file = %s WHERE id = %s",
(avatar_file, conversation_id),
)
conn.commit()
finally:
conn.close()
def update_conversation_name(conversation_id: str, name: str):
"""Update the name of a conversation."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"UPDATE conversations SET name = %s WHERE id = %s",
(name, conversation_id),
)
conn.commit()
finally:
conn.close()
def is_conversation_member(conversation_id: str, user_id: str) -> bool:
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"SELECT 1 FROM conversation_members WHERE conversation_id = %s AND user_id = %s",
(conversation_id, user_id),
)
return cursor.fetchone() is not None
finally:
conn.close()
def list_user_conversations(user_id: str) -> list[dict]:
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT c.id, c.created_at, c.name, c.created_by, c.avatar_file FROM conversations c "
"JOIN conversation_members cm ON c.id = cm.conversation_id "
"WHERE cm.user_id = %s ORDER BY c.created_at DESC",
(user_id,),
)
convs = cursor.fetchall()
for conv in convs:
cursor.execute(
"SELECT u.id AS user_id, u.username, u.email FROM conversation_members cm "
"JOIN users u ON cm.user_id = u.id "
"WHERE cm.conversation_id = %s",
(conv["id"],),
)
conv["members"] = cursor.fetchall()
return convs
finally:
conn.close()
# --- Group Invitations ---
def create_invitation(conversation_id: str, user_id: str, invited_by: str):
"""Create a pending group invitation."""
conn = get_connection()
try:
cursor = conn.cursor()
inv_id = generate_uuid()
cursor.execute(
"INSERT IGNORE INTO group_invitations (id, conversation_id, user_id, invited_by) "
"VALUES (%s, %s, %s, %s)",
(inv_id, conversation_id, user_id, invited_by),
)
conn.commit()
finally:
conn.close()
def get_pending_invitations(user_id: str) -> list[dict]:
"""Get all pending invitations for a user, joined with conversation and inviter info."""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT gi.id, gi.conversation_id, gi.invited_by, gi.created_at, "
"c.name AS conversation_name, u.username AS invited_by_username "
"FROM group_invitations gi "
"JOIN conversations c ON gi.conversation_id = c.id "
"JOIN users u ON gi.invited_by = u.id "
"WHERE gi.user_id = %s "
"ORDER BY gi.created_at DESC",
(user_id,),
)
return cursor.fetchall()
finally:
conn.close()
def delete_invitation(conversation_id: str, user_id: str):
"""Delete a pending invitation."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"DELETE FROM group_invitations WHERE conversation_id = %s AND user_id = %s",
(conversation_id, user_id),
)
conn.commit()
finally:
conn.close()
def has_pending_invitation(conversation_id: str, user_id: str) -> bool:
"""Check if a user has a pending invitation for a conversation."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"SELECT 1 FROM group_invitations WHERE conversation_id = %s AND user_id = %s",
(conversation_id, user_id),
)
return cursor.fetchone() is not None
finally:
conn.close()
# --- Messages ---
def store_message(
conversation_id: str,
sender_id: str,
ratchet_header: bytes,
recipients: list[dict],
x3dh_header: bytes | None = None,
sender_chain_id: bytes | None = None,
sender_chain_n: int | None = None,
image_file_id: str | None = None,
sender_device_id: str | None = None,
) -> str:
"""Store an encrypted message with per-recipient ciphertext.
recipients: [{user_id, encrypted_content (bytes), nonce (bytes),
device_id (str, optional), ratchet_header (bytes, optional),
x3dh_header (bytes, optional)}]
"""
conn = get_connection()
try:
cursor = conn.cursor()
msg_id = generate_uuid()
cursor.execute(
"INSERT INTO messages (id, conversation_id, sender_id, sender_device_id, "
"ratchet_header, x3dh_header, sender_chain_id, sender_chain_n, image_file_id) "
"VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)",
(msg_id, conversation_id, sender_id, sender_device_id, ratchet_header,
x3dh_header, sender_chain_id, sender_chain_n, image_file_id),
)
for r in recipients:
device_id = r.get("device_id", SELF_DEVICE_ID)
cursor.execute(
"INSERT INTO message_recipients (message_id, user_id, device_id, "
"encrypted_content, nonce, ratchet_header, x3dh_header) "
"VALUES (%s, %s, %s, %s, %s, %s, %s)",
(msg_id, r["user_id"], device_id, r["encrypted_content"], r["nonce"],
r.get("ratchet_header"), r.get("x3dh_header")),
)
conn.commit()
return msg_id
finally:
conn.close()
def get_messages(conversation_id: str, user_id: str, limit: int = 50, offset: int = 0,
device_id: str | None = None) -> list[dict]:
"""Get messages for a user in a conversation, JOINing their per-recipient ciphertext.
If device_id is set, returns rows where mr.device_id matches OR is the sentinel
(self-encrypted / legacy). This ensures both device-specific and self-encrypted
copies are returned.
"""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
if device_id:
cursor.execute(
"SELECT m.id, m.conversation_id, m.sender_id, m.sender_device_id, "
"m.ratchet_header, m.x3dh_header, "
"m.sender_chain_id, m.sender_chain_n, m.created_at, m.deleted_at, m.image_file_id, "
"mr.encrypted_content, mr.nonce, mr.device_id AS mr_device_id, "
"mr.ratchet_header AS mr_ratchet_header, mr.x3dh_header AS mr_x3dh_header "
"FROM messages m "
"JOIN message_recipients mr ON m.id = mr.message_id AND mr.user_id = %s "
" AND (mr.device_id = %s OR mr.device_id = %s) "
"JOIN conversation_members cm ON cm.conversation_id = m.conversation_id AND cm.user_id = %s "
"WHERE m.conversation_id = %s AND (cm.joined_at IS NULL OR m.created_at >= cm.joined_at) "
"ORDER BY m.created_at DESC LIMIT %s OFFSET %s",
(user_id, device_id, SELF_DEVICE_ID, user_id, conversation_id, limit, offset),
)
else:
cursor.execute(
"SELECT m.id, m.conversation_id, m.sender_id, m.sender_device_id, "
"m.ratchet_header, m.x3dh_header, "
"m.sender_chain_id, m.sender_chain_n, m.created_at, m.deleted_at, m.image_file_id, "
"mr.encrypted_content, mr.nonce, mr.device_id AS mr_device_id, "
"mr.ratchet_header AS mr_ratchet_header, mr.x3dh_header AS mr_x3dh_header "
"FROM messages m "
"JOIN message_recipients mr ON m.id = mr.message_id AND mr.user_id = %s "
"JOIN conversation_members cm ON cm.conversation_id = m.conversation_id AND cm.user_id = %s "
"WHERE m.conversation_id = %s AND (cm.joined_at IS NULL OR m.created_at >= cm.joined_at) "
"ORDER BY m.created_at DESC LIMIT %s OFFSET %s",
(user_id, user_id, conversation_id, limit, offset),
)
return cursor.fetchall()
finally:
conn.close()
def get_message_conversation(message_id: str) -> str | None:
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute("SELECT conversation_id FROM messages WHERE id = %s", (message_id,))
row = cursor.fetchone()
return row[0] if row else None
finally:
conn.close()
def get_message_sender(message_id: str) -> str | None:
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute("SELECT sender_id FROM messages WHERE id = %s", (message_id,))
row = cursor.fetchone()
return row[0] if row else None
finally:
conn.close()
# --- Group Sender Keys ---
def store_sender_key(conversation_id: str, sender_id: str, chain_id: bytes,
device_id: str | None = None):
"""Store or update a sender key chain ID for a group member's device."""
conn = get_connection()
try:
cursor = conn.cursor()
dev = device_id or SELF_DEVICE_ID
cursor.execute(
"REPLACE INTO group_sender_keys (conversation_id, sender_id, device_id, chain_id) "
"VALUES (%s, %s, %s, %s)",
(conversation_id, sender_id, dev, chain_id),
)
conn.commit()
finally:
conn.close()
def get_sender_key(conversation_id: str, sender_id: str,
device_id: str | None = None) -> dict | None:
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
dev = device_id or SELF_DEVICE_ID
cursor.execute(
"SELECT chain_id, created_at FROM group_sender_keys "
"WHERE conversation_id = %s AND sender_id = %s AND device_id = %s",
(conversation_id, sender_id, dev),
)
return cursor.fetchone()
finally:
conn.close()
# --- Read Receipts ---
def mark_messages_read(conversation_id: str, user_id: str, message_ids: list[str]):
if not message_ids:
return
conn = get_connection()
try:
cursor = conn.cursor()
for mid in message_ids:
cursor.execute(
"INSERT IGNORE INTO message_reads (message_id, user_id) VALUES (%s, %s)",
(mid, user_id),
)
conn.commit()
finally:
conn.close()
def get_unread_counts(user_id: str) -> dict[str, int]:
"""Return {conversation_id: unread_count} for all conversations the user is in."""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT m.conversation_id, COUNT(*) AS cnt "
"FROM messages m "
"JOIN message_recipients mr ON mr.message_id = m.id AND mr.user_id = %s "
"LEFT JOIN message_reads mrd ON mrd.message_id = m.id AND mrd.user_id = %s "
"WHERE m.sender_id != %s AND m.deleted_at IS NULL AND mrd.message_id IS NULL "
"GROUP BY m.conversation_id",
(user_id, user_id, user_id),
)
return {row["conversation_id"]: row["cnt"] for row in cursor.fetchall()}
finally:
conn.close()
def get_message_read_status(message_ids: list[str]) -> dict:
if not message_ids:
return {}
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
placeholders = ",".join(["%s"] * len(message_ids))
cursor.execute(
f"SELECT mr.message_id, mr.user_id, mr.read_at "
f"FROM message_reads mr "
f"WHERE mr.message_id IN ({placeholders})",
tuple(message_ids),
)
result = {}
for row in cursor.fetchall():
mid = row["message_id"]
if mid not in result:
result[mid] = []
result[mid].append({
"user_id": row["user_id"],
"read_at": row["read_at"].isoformat() if hasattr(row["read_at"], "isoformat") else str(row["read_at"]),
})
return result
finally:
conn.close()
# --- Delete ---
def soft_delete_message(message_id: str, sender_id: str) -> dict | None:
"""Soft-delete a message if sender matches. Returns {'image_file_id': ...} or None."""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT sender_id, image_file_id FROM messages WHERE id = %s AND deleted_at IS NULL",
(message_id,),
)
row = cursor.fetchone()
if not row or row["sender_id"] != sender_id:
return None
cursor.execute(
"UPDATE messages SET deleted_at = NOW() WHERE id = %s",
(message_id,),
)
# Clear per-recipient ciphertext
cursor.execute(
"UPDATE message_recipients SET encrypted_content = %s WHERE message_id = %s",
(b"", message_id),
)
conn.commit()
return {"image_file_id": row.get("image_file_id")}
finally:
conn.close()
def set_message_image_file_id(message_id: str, file_id: str):
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"UPDATE messages SET image_file_id = %s WHERE id = %s",
(file_id, message_id),
)
conn.commit()
finally:
conn.close()
# --- Image Uploads ---
def create_image_upload(file_id: str, conversation_id: str, uploader_id: str, file_size: int):
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"INSERT INTO image_uploads (file_id, conversation_id, uploader_id, file_size) "
"VALUES (%s, %s, %s, %s)",
(file_id, conversation_id, uploader_id, file_size),
)
conn.commit()
finally:
conn.close()
def complete_image_upload(file_id: str):
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"UPDATE image_uploads SET completed = TRUE WHERE file_id = %s",
(file_id,),
)
conn.commit()
finally:
conn.close()
def get_image_upload(file_id: str) -> dict | None:
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT file_id, conversation_id, uploader_id, file_size, completed, created_at "
"FROM image_uploads WHERE file_id = %s",
(file_id,),
)
return cursor.fetchone()
finally:
conn.close()
def delete_image_upload(file_id: str):
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute("DELETE FROM image_uploads WHERE file_id = %s", (file_id,))
conn.commit()
finally:
conn.close()
# --- User Profiles ---
def create_default_profile(user_id: str):
"""Create a default profile for a new user."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"INSERT IGNORE INTO user_profiles (user_id) VALUES (%s)",
(user_id,),
)
conn.commit()
finally:
conn.close()
def get_user_profile(user_id: str, viewer_id: str | None = None) -> dict | None:
"""Get user profile joined with user info. Respects visibility if viewer is different user."""
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT u.id AS user_id, u.username, u.email, u.created_at, "
"p.phone, p.phone_visible, p.email_visible, p.location, "
"p.location_visible, p.avatar_file, p.updated_at "
"FROM users u LEFT JOIN user_profiles p ON u.id = p.user_id "
"WHERE u.id = %s",
(user_id,),
)
row = cursor.fetchone()
if not row:
return None
# If viewing someone else's profile, apply visibility rules
if viewer_id and viewer_id != user_id:
if not row.get("email_visible"):
row["email"] = None
if not row.get("phone_visible"):
row["phone"] = None
if not row.get("location_visible"):
row["location"] = None
return row
finally:
conn.close()
def update_user_profile(user_id: str, **fields):
"""Upsert user profile fields. Allowed: phone, phone_visible, email_visible,
location, location_visible, avatar_file."""
allowed = {"phone", "phone_visible", "email_visible", "location",
"location_visible", "avatar_file"}
filtered = {k: v for k, v in fields.items() if k in allowed}
if not filtered:
return
conn = get_connection()
try:
cursor = conn.cursor()
# Upsert: insert default then update
cursor.execute(
"INSERT IGNORE INTO user_profiles (user_id) VALUES (%s)",
(user_id,),
)
set_clause = ", ".join(f"{k} = %s" for k in filtered)
values = list(filtered.values()) + [user_id]
cursor.execute(
f"UPDATE user_profiles SET {set_clause} WHERE user_id = %s",
values,
)
conn.commit()
finally:
conn.close()
def batch_reencrypt_messages(user_id: str, updates: list[dict]):
"""Batch re-encrypt message_recipients rows with self-encryption key data.
Each update: {message_id, encrypted_content (bytes), nonce (bytes)}.
Sets ratchet_header to '{"self":true}' and clears x3dh_header.
Only updates rows belonging to user_id with sentinel device_id (self-encrypted copies).
"""
if not updates:
return
conn = get_connection()
try:
cursor = conn.cursor()
self_header = b'{"self":true}'
for u in updates:
cursor.execute(
"UPDATE message_recipients "
"SET encrypted_content = %s, nonce = %s, ratchet_header = %s, x3dh_header = NULL "
"WHERE message_id = %s AND user_id = %s AND device_id = %s",
(u["encrypted_content"], u["nonce"], self_header, u["message_id"],
user_id, SELF_DEVICE_ID),
)
conn.commit()
finally:
conn.close()
# --- Phantom Users ---
def create_phantom_user(email: str) -> dict:
"""Create a phantom user with valid crypto keys for X3DH.
Phantom users have rsa_public_key = 'PHANTOM' as a marker.
Returns user dict: {id, username, email, identity_key}.
"""
username = email.split("@")[0]
user_id = generate_uuid()
# Generate real crypto keys so X3DH works on the client side
ik_private, ik_public = generate_identity_keypair()
ik_public_bytes = serialize_ed25519_public(ik_public)
spk = generate_signed_prekey(ik_private)
spk_pub_bytes = serialize_x25519_public(spk["public"])
spk_sig = spk["signature"]
opks = generate_one_time_prekeys(count=5)
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"INSERT INTO users (id, username, email, rsa_public_key, identity_key) "
"VALUES (%s, %s, %s, %s, %s)",
(user_id, username, email, "PHANTOM", ik_public_bytes),
)
cursor.execute(
"INSERT INTO signed_prekeys (id, user_id, public_key, signature) VALUES (%s, %s, %s, %s)",
(spk["id"], user_id, spk_pub_bytes, spk_sig),
)
for opk in opks:
cursor.execute(
"INSERT INTO one_time_prekeys (id, user_id, public_key) VALUES (%s, %s, %s)",
(opk["id"], user_id, serialize_x25519_public(opk["public"])),
)
conn.commit()
return {"id": user_id, "username": username, "email": email, "identity_key": ik_public_bytes}
finally:
conn.close()
def is_phantom_user(user_id: str) -> bool:
"""Check if a user is a phantom (rsa_public_key == 'PHANTOM')."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute("SELECT rsa_public_key FROM users WHERE id = %s", (user_id,))
row = cursor.fetchone()
return row is not None and row[0] == "PHANTOM"
finally:
conn.close()
def delete_phantom_user(user_id: str):
"""Delete a phantom user. CASCADE removes signed_prekeys, one_time_prekeys,
conversation_members, message_recipients, etc."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"DELETE FROM users WHERE id = %s AND rsa_public_key = %s",
(user_id, "PHANTOM"),
)
conn.commit()
finally:
conn.close()
def upgrade_phantom_user(phantom_id: str, username: str, rsa_public_key_pem: str,
identity_key: bytes) -> str | None:
"""Upgrade a phantom user to a real user in-place.
Preserves user_id and all FK references (conversation_members, group_invitations, etc.).
Deletes phantom's server-generated prekeys (real user will upload own on first login).
Returns phantom_id as the new user_id, or None if phantom no longer exists.
"""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"UPDATE users SET username = %s, rsa_public_key = %s, identity_key = %s "
"WHERE id = %s AND rsa_public_key = 'PHANTOM'",
(username, rsa_public_key_pem, identity_key, phantom_id),
)
if cursor.rowcount == 0:
conn.rollback()
return None
# Remove phantom's server-generated crypto keys — real user uploads own
cursor.execute("DELETE FROM signed_prekeys WHERE user_id = %s", (phantom_id,))
cursor.execute("DELETE FROM one_time_prekeys WHERE user_id = %s", (phantom_id,))
conn.commit()
return phantom_id
finally:
conn.close()
def get_all_phantom_user_ids() -> set[str]:
"""Return set of all phantom user IDs (for server startup cache)."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute("SELECT id FROM users WHERE rsa_public_key = %s", ("PHANTOM",))
return {row[0] for row in cursor.fetchall()}
finally:
conn.close()
def cleanup_stale_phantoms(max_age_days: int = 30) -> int:
"""Delete phantom users older than max_age_days with no active conversations with real users."""
conn = get_connection()
try:
cursor = conn.cursor()
# Two-step: SELECT ids first, then DELETE.
# MySQL error 1093: can't DELETE from table referenced in subquery.
cursor.execute("""
SELECT u.id FROM users u
WHERE u.rsa_public_key = 'PHANTOM'
AND u.created_at < DATE_SUB(NOW(), INTERVAL %s DAY)
AND NOT EXISTS (
SELECT 1 FROM conversation_members cm1
JOIN conversation_members cm2 ON cm1.conversation_id = cm2.conversation_id
JOIN users u2 ON cm2.user_id = u2.id
WHERE cm1.user_id = u.id
AND u2.rsa_public_key != 'PHANTOM'
)
""", (max_age_days,))
ids = [row[0] for row in cursor.fetchall()]
if not ids:
return 0
cursor.execute(
"DELETE FROM users WHERE id IN (%s)" % ",".join(["%s"] * len(ids)),
ids,
)
deleted = cursor.rowcount
conn.commit()
return deleted
finally:
conn.close()
def remove_conversation_member_atomic(conversation_id: str, user_id: str) -> bool:
"""Remove member and return True if actually removed (row existed). M6 TOCTOU fix."""
conn = get_connection()
try:
cursor = conn.cursor()
cursor.execute(
"DELETE FROM conversation_members WHERE conversation_id = %s AND user_id = %s",
(conversation_id, user_id),
)
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
def get_stale_uploads(max_age_seconds: int = 3600) -> list[dict]:
conn = get_connection()
try:
cursor = conn.cursor(dictionary=True)
cursor.execute(
"SELECT file_id FROM image_uploads "
"WHERE completed = FALSE AND created_at < DATE_SUB(NOW(), INTERVAL %s SECOND)",
(max_age_seconds,),
)
return cursor.fetchall()
finally:
conn.close()