replypier/backend/tests/core/test_encryption.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import base64
import sys
from pathlib import Path
from uuid import uuid4

import pytest
from cryptography.fernet import Fernet
from sqlalchemy import text

from app.core.config import settings
from app.core.encryption import (
    CURRENT_KEY_ID,
    ENCRYPTED_PREFIX,
    ENCRYPTION_KEY_ENV,
    JSON_MARKER,
    SEPARATOR,
    EncryptedText,
    decrypt_json,
    decrypt_value,
    encrypt_json,
    encrypt_value,
    is_encrypted,
    load_key,
)
from app.database import Base
from app.models.chat_history import ChatHistory
from app.models.session_to_agent import SessionToAgent, SessionStatus

# Where an installed extension's models live (see app.channels.app_hooks)
EXTENSION_PACKAGE = "app.enterprise."


def test_round_trip():
    token = encrypt_value("hello customer")
    assert token.startswith(f"{ENCRYPTED_PREFIX}{SEPARATOR}{CURRENT_KEY_ID}{SEPARATOR}")
    assert "hello customer" not in token
    assert decrypt_value(token) == "hello customer"


def test_encrypt_is_idempotent():
    token = encrypt_value("hello")
    assert encrypt_value(token) == token


def test_plaintext_passes_through():
    """Rows the backfill has not reached yet must keep reading."""
    assert decrypt_value("not yet encrypted") == "not yet encrypted"


def test_none_and_empty():
    assert encrypt_value(None) is None
    assert decrypt_value(None) is None
    assert decrypt_value("") == ""
    assert decrypt_value(encrypt_value("")) == ""


def test_unknown_key_id_returns_none():
    token = encrypt_value("secret").replace(CURRENT_KEY_ID, "v99", 1)
    assert decrypt_value(token) is None


def test_corrupt_token_returns_none():
    assert decrypt_value(f"{ENCRYPTED_PREFIX}{SEPARATOR}{CURRENT_KEY_ID}{SEPARATOR}nope") is None


def test_malformed_ciphertext_returns_none():
    assert decrypt_value(f"{ENCRYPTED_PREFIX}{SEPARATOR}{CURRENT_KEY_ID}") is None


def test_is_encrypted():
    assert is_encrypted(encrypt_value("x"))
    assert not is_encrypted("plain")
    assert not is_encrypted(None)


def test_json_round_trip():
    memory = {"runs": [{"content": "my secret transcript"}]}
    blob = encrypt_json(memory)
    assert list(blob) == [JSON_MARKER]
    assert "secret" not in str(blob)
    assert decrypt_json(blob) == memory


def test_json_plaintext_passes_through():
    legacy = {"runs": [{"content": "written before encryption"}]}
    assert decrypt_json(legacy) == legacy
    assert decrypt_json(None) is None
    assert encrypt_json(None) is None


def test_json_encrypt_is_idempotent():
    blob = encrypt_json({"a": 1})
    assert encrypt_json(blob) == blob


def test_undecryptable_json_returns_none():
    assert decrypt_json({JSON_MARKER: f"{ENCRYPTED_PREFIX}{SEPARATOR}{CURRENT_KEY_ID}{SEPARATOR}x"}) is None


def test_backfill_script_covers_every_encrypted_column():
    """The backfill lives outside the app tree and cannot import the models, so
    nothing but this test stops a newly encrypted column from shipping with its
    existing rows silently left in plaintext."""
    sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "scripts"))
    try:
        from encrypt_chat_at_rest import TARGETS
    finally:
        sys.path.pop(0)

    # Only the open-source models: an installed extension (app.enterprise)
    # owns its tables and any backfill they need.
    declared = {
        (mapper.local_table.name, column.name)
        for mapper in Base.registry.mappers
        if not mapper.class_.__module__.startswith(EXTENSION_PACKAGE)
        for column in mapper.local_table.columns
        if isinstance(column.type, EncryptedText)
    }
    covered = {
        (table, column)
        for table, _pk, columns, _kind in TARGETS
        for column in columns
    }
    assert declared <= covered, f"not backfilled: {declared - covered}"


def test_chat_message_is_ciphertext_in_the_database(db, test_organization_id):
    """The ORM must round-trip, and the stored bytes must not be readable."""
    message = ChatHistory(
        organization_id=test_organization_id,
        message="card number is 4111 1111 1111 1111",
        message_type="user",
    )
    db.add(message)
    db.commit()

    # Raw SQL: selecting the typed column would decrypt it again on the way out.
    raw = db.execute(
        text("SELECT message FROM chat_history WHERE id = :id"), {"id": message.id}
    ).scalar()
    assert is_encrypted(raw)
    assert "4111" not in raw

    db.expire_all()
    assert db.get(ChatHistory, message.id).message == "card number is 4111 1111 1111 1111"


def test_session_summary_is_ciphertext_in_the_database(db, test_organization_id):
    session = SessionToAgent(
        session_id=uuid4(),
        organization_id=test_organization_id,
        customer_id=uuid4(),
        status=SessionStatus.OPEN,
        ticket_summary="customer threatened to churn",
        transfer_description="needs a human",
    )
    db.add(session)
    db.commit()

    raw = db.execute(
        text("SELECT ticket_summary FROM session_to_agents "
             "WHERE ticket_summary IS NOT NULL")
    ).scalar()
    assert is_encrypted(raw)
    assert "churn" not in raw

    db.expire_all()
    reloaded = db.get(SessionToAgent, session.session_id)
    assert reloaded.ticket_summary == "customer threatened to churn"
    assert reloaded.transfer_description == "needs a human"


def test_legacy_plaintext_row_still_reads(db, test_organization_id):
    """A row written before this change has no prefix and must not break reads."""
    # Raw SQL on purpose: going through the ORM would encrypt it.
    db.execute(
        text("INSERT INTO chat_history (organization_id, message, message_type) "
             "VALUES (:org, :message, :message_type)"),
        {"org": str(test_organization_id), "message": "written before encryption",
         "message_type": "user"},
    )
    db.commit()

    stored = db.query(ChatHistory).filter(
        ChatHistory.message_type == "user"
    ).order_by(ChatHistory.id.desc()).first()
    assert stored.message == "written before encryption"


def _double_encoded(key: bytes) -> str:
    """The wire format ENCRYPTION_KEY uses: a Fernet key, base64-encoded again."""
    return base64.b64encode(key).decode()


def test_load_key_accepts_a_double_encoded_key(monkeypatch):
    monkeypatch.setattr(settings, "ENVIRONMENT", "production")
    key = Fernet.generate_key()
    monkeypatch.setenv(ENCRYPTION_KEY_ENV, _double_encoded(key))
    assert load_key() == key


def test_load_key_rejects_base64_that_is_not_a_fernet_key(monkeypatch):
    """Encoding the key once instead of twice decodes cleanly but is 32 raw bytes,
    not Fernet key material — it must fail here, not later inside Fernet()."""
    monkeypatch.setattr(settings, "ENVIRONMENT", "production")
    monkeypatch.setenv(ENCRYPTION_KEY_ENV, base64.b64encode(b"x" * 32).decode())
    with pytest.raises(RuntimeError, match="not a usable Fernet key"):
        load_key()


def test_load_key_rejects_the_env_example_placeholder(monkeypatch):
    monkeypatch.setattr(settings, "ENVIRONMENT", "production")
    monkeypatch.setenv(ENCRYPTION_KEY_ENV, "your_fernet_encryption_key_here")
    with pytest.raises(RuntimeError, match="not a usable Fernet key"):
        load_key()


def test_load_key_is_fatal_when_unset_outside_development(monkeypatch):
    monkeypatch.setattr(settings, "ENVIRONMENT", "production")
    monkeypatch.delenv(ENCRYPTION_KEY_ENV, raising=False)
    with pytest.raises(RuntimeError, match="is not set"):
        load_key()


def test_load_key_falls_back_to_a_throwaway_key_in_development(monkeypatch):
    """An unusable key must not stop local development — a generated one is fine
    there, since nothing durable is encrypted under it."""
    monkeypatch.setattr(settings, "ENVIRONMENT", "development")
    monkeypatch.setenv(ENCRYPTION_KEY_ENV, "your_fernet_encryption_key_here")
    Fernet(load_key())  # usable, and no raise