"""Tests for the DeDRM crypto primitives (alfcrypto)."""

import hashlib

import pytest

import dedrm_test_utils as U

# Load via the `dedrm` package so we exercise the same module object that
# topazextract/mobidedrm import through `from .alfcrypto import ...`, rather
# than a second top-level copy.
alf = U.load("alfcrypto", package="dedrm")


def test_pc1_roundtrip():
    key = bytes(range(16))
    plaintext = b"Hello, Topaz DRM world! 12345678"
    encrypted = alf.Pukall_Cipher().PC1(key, plaintext, decryption=False)
    assert encrypted != plaintext
    assert alf.Pukall_Cipher().PC1(key, encrypted, decryption=True) == plaintext


def test_pc1_rejects_bad_key_length():
    # Pin the specific guard, not just "some exception".
    with pytest.raises(Exception, match="Bad key length"):
        alf.Pukall_Cipher().PC1(b"\x00" * 8, b"data............", decryption=False)


def _topaz_encrypt(data, ctx):
    """Inverse of Topaz_Cipher.decrypt, used to drive a round-trip test."""
    ctx1, ctx2 = ctx
    out = bytearray()
    for plain in data:
        out.append((plain ^ ((ctx1 >> 3) & 0xFF) ^ ((ctx2 << 3) & 0xFF)) & 0xFF)
        ctx2 = ctx1
        ctx1 = (((ctx1 >> 2) * (ctx1 >> 7)) & 0xFFFFFFFF) ^ ((plain * plain * 0x0F902007) & 0xFFFFFFFF)
    return bytes(out)


def test_topaz_cipher_roundtrip():
    plaintext = b"topaz plaintext bytes"
    cipher = alf.Topaz_Cipher()
    ctx = cipher.ctx_init(b"mykey123")
    ciphertext = _topaz_encrypt(plaintext, list(ctx))
    recovered = cipher.decrypt(ciphertext, list(ctx)).encode("latin-1")
    assert recovered == plaintext


def test_topaz_ctx_init_golden():
    # Independent golden vector pins the key-schedule math; a bare determinism
    # check (ctx_init(x) == ctx_init(x)) would pass for any implementation.
    assert alf.Topaz_Cipher().ctx_init(b"abc") == [426181496, 3966955212]


def test_topaz_decrypt_golden():
    # Golden oracle independent of the test's own inverse helper: a fixed
    # ciphertext + key must always decrypt to exactly these bytes. This catches
    # a systematic state-update bug that the round-trip test cannot.
    ctx = alf.Topaz_Cipher().ctx_init(b"topaz-key")
    out = alf.Topaz_Cipher().decrypt(bytes(range(32)), list(ctx)).encode("latin-1")
    assert out == bytes.fromhex(
        "e64653fbb4c4ad06981b51a0e53cd17ad9c7bc72932626ed35b935791f4fd769"
    )


def test_pbkdf2_matches_hashlib():
    out = alf.KeyIVGen().pbkdf2(b"password", b"salt", 100, 32)
    assert out == hashlib.pbkdf2_hmac("sha1", b"password", b"salt", 100, 32)