"""Tests for mobidedrm helpers."""

import pytest

import dedrm_test_utils as U

md = U.load("mobidedrm", package="dedrm")


def test_pc1_roundtrip():
    key = bytes(range(16))
    plaintext = b"cookie-data-1234"
    ciphertext = md.PC1(key, plaintext, decryption=False)
    assert ciphertext != plaintext
    assert md.PC1(key, ciphertext, decryption=True) == plaintext


def test_pc1_rejects_bad_key_length():
    with pytest.raises(Exception, match="Bad key length"):
        md.PC1(b"\x00" * 8, b"data............", decryption=False)


def test_trailing_data_entries():
    buf = b"XXXX" + bytes([0x82])
    # bit 1 set -> one size entry encoded in the final byte (0x82 -> 2)
    assert md.getSizeOfTrailingDataEntries(buf, len(buf), 0b10) == 2
    # no flags -> no trailing data
    assert md.getSizeOfTrailingDataEntries(buf, len(buf), 0) == 0
    # bit 0 set -> multibyte entry: (last_byte & 3) + 1
    assert md.getSizeOfTrailingDataEntries(b"ABCDE", 5, 1) == 2


def test_checksum_pid_str_contract():
    # mobidedrm's call sites pass and expect ``str``; the shared helper must
    # preserve that (unlike the bytes contract used by kgenpids/kindlepid).
    out = md.checksumPid("12345678")
    assert out == "12345678EL"
    assert isinstance(out, str)