// SPDX-FileCopyrightText: © 2026 Jeffrey C. Ollie <jeff@ocjtech.us>
// SPDX-License-Identifier: MIT

//! Mobipocket key derivation.
//!
//! Two schemes exist. The legacy scheme (`encryption_type` 1) stores the book
//! key in record zero, obfuscated with a fixed vector, so nothing is needed
//! from the caller. The DRM scheme (`encryption_type` 2) stores the key in a
//! voucher that opens either with a fixed "reader" key or with a key derived
//! from the reading device's serial number.
//!
//! Only key *derivation* lives here; the cipher itself is in `pc1.zig`.

const std = @import("std");
const pc1 = @import("pc1.zig");

pub const Error = error{
    InvalidSerial,
    KeyNotFound,
};

/// Opens vouchers that are not bound to a device.
pub const internal_reader_key: pc1.Key = .{
    0x72, 0x38, 0x33, 0xb0, 0xb4, 0xf2, 0xe3, 0xca,
    0xdf, 0x09, 0x01, 0xd6, 0xe2, 0xe0, 0x3f, 0x96,
};

/// Obfuscates the book key stored in record zero by the legacy scheme.
pub const internal_reader_key_v1: pc1.Key = "QDCVEPMU675RUBSZ".*;

pub const pid_size = 10;
pub const Pid = [pid_size]u8;

/// Serial numbers come in these two lengths; nothing else is accepted.
pub const serial_size = 16;
pub const serial_long_size = 40;

/// The alphabet PIDs are spelled in. Note the missing 'O' and '0'.
const alphabet = "ABCDEFGHIJKLMNPQRSTUVWXYZ123456789";

const crc_table: [256]u32 = blk: {
    @setEvalBranchQuota(20000);
    var table: [256]u32 = undefined;
    for (&table, 0..) |*entry, i| {
        var c: u32 = @intCast(i);
        for (0..8) |_| c = if (c & 1 != 0) 0xEDB88320 ^ (c >> 1) else c >> 1;
        entry.* = c;
    }
    break :blk table;
};

/// The reflected CRC-32 register (polynomial 0xEDB88320) starting from zero
/// with no final inversion. This is *not* the usual CRC-32 value: the standard
/// one starts the register at 0xFFFFFFFF and inverts at the end. Mobipocket
/// uses the bare register, so the two disagree for every input.
fn crcRegister(data: []const u8) u32 {
    var c: u32 = 0;
    for (data) |b| c = crc_table[(c ^ b) & 0xff] ^ (c >> 8);
    return c;
}

/// The two trailing characters of a PID, which check the first eight.
fn pidChecksum(pid: Pid) [2]u8 {
    var crc = crcRegister(pid[0 .. pid_size - 2]);
    crc ^= crc >> 16;
    var out: [2]u8 = undefined;
    for (&out) |*o| {
        const b: u8 = @truncate(crc);
        const pos = (b / alphabet.len) ^ (b % alphabet.len);
        o.* = alphabet[pos % alphabet.len];
        crc >>= 8;
    }
    return out;
}

/// Derive the PID for a device from its serial number.
pub fn pidFromSerial(serial: []const u8) Error!Pid {
    if (serial.len != serial_size and serial.len != serial_long_size) return error.InvalidSerial;

    var pid: Pid = @splat(0);
    var scratch: [pid_size - 2]u8 = @splat(0);
    var len = scratch.len;
    if (serial.len == serial_size) {
        // Short serials leave a literal '*' in the last body position.
        len -= 1;
        pid[len] = '*';
    }

    const crc = crcRegister(serial);
    for (serial, 0..) |c, i| scratch[i % len] ^= c;
    for (0..len) |i| {
        const shift: u5 = @intCast(24 - 8 * (i & 3));
        scratch[i] ^= @as(u8, @truncate(crc >> shift));
        const v = scratch[i];
        // Caps at 32, so it always lands inside the 34 character alphabet.
        const pos = (v >> 7) + ((v >> 5 & 3) ^ (v & 0x1f));
        pid[i] = alphabet[pos];
    }

    const checksum = pidChecksum(pid);
    pid[pid_size - 2] = checksum[0];
    pid[pid_size - 1] = checksum[1];
    return pid;
}

pub fn pidIsValid(pid: Pid) bool {
    return std.mem.eql(u8, &pidChecksum(pid), pid[pid_size - 2 ..]);
}

/// The cipher key a PID unlocks vouchers with.
pub fn pidKey(pid: Pid) pc1.Key {
    var key: pc1.Key = @splat(0);
    @memcpy(key[0 .. pid_size - 2], pid[0 .. pid_size - 2]);
    pc1.encrypt(&key, &key, internal_reader_key);
    return key;
}

/// Vouchers record which key opens them as a one byte sum, so the wrong key can
/// be skipped without doing the work of decrypting the cookie.
pub fn keyChecksum(key: pc1.Key) u8 {
    var sum: usize = 0;
    for (key) |b| sum += b;
    return @truncate(sum);
}

pub const voucher_size = 48;
pub const cookie_size = 32;

/// One entry in record zero's DRM section.
pub const Voucher = struct {
    verification: u32,
    size: u32,
    type: u32,
    checksum: u8,
    cookie: [cookie_size]u8,

    pub fn parse(bytes: *const [voucher_size]u8) Voucher {
        return .{
            .verification = std.mem.readInt(u32, bytes[0..4], .big),
            .size = std.mem.readInt(u32, bytes[4..8], .big),
            .type = std.mem.readInt(u32, bytes[8..12], .big),
            .checksum = bytes[12],
            // Three bytes of padding sit between the checksum and the cookie.
            .cookie = bytes[16..48].*,
        };
    }
};

/// What a successfully opened voucher yields.
pub const Unlocked = struct {
    key: pc1.Key,
    /// Validity window, in minutes since the Unix epoch. `valid_from` of zero
    /// or `valid_to` of 0xFFFFFFFF means the voucher never expires. This
    /// library reports the window but does not enforce it.
    valid_from: u32,
    valid_to: u32,

    pub fn expiresAt(self: Unlocked) ?u32 {
        if (self.valid_from == 0 or self.valid_to == 0xffff_ffff) return null;
        return self.valid_to;
    }
};

/// Key types, as stored in the low five bits of the cookie's flags.
pub const key_type_device = 1;
pub const key_type_internal = 3;

/// Try `key` against one voucher. Returns null if it is not the right key.
pub fn openVoucher(voucher: Voucher, key: pc1.Key, key_type: u8) ?Unlocked {
    if (voucher.checksum != keyChecksum(key)) return null;

    var cookie: [cookie_size]u8 = undefined;
    pc1.decrypt(&cookie, &voucher.cookie, key);

    if (std.mem.readInt(u32, cookie[0..4], .big) != voucher.verification) return null;
    const flags = std.mem.readInt(u32, cookie[4..8], .big);
    if (flags & 0x1f != key_type) return null;

    return .{
        .key = cookie[8..24].*,
        .valid_to = std.mem.readInt(u32, cookie[24..28], .big),
        .valid_from = std.mem.readInt(u32, cookie[28..32], .big),
    };
}

/// Build a voucher that `key` opens. Used by the tests, and the inverse of
/// `openVoucher`.
pub fn sealVoucher(book_key: pc1.Key, key: pc1.Key, key_type: u8, verification: u32) [voucher_size]u8 {
    var cookie: [cookie_size]u8 = @splat(0);
    std.mem.writeInt(u32, cookie[0..4], verification, .big);
    std.mem.writeInt(u32, cookie[4..8], key_type, .big);
    @memcpy(cookie[8..24], &book_key);
    std.mem.writeInt(u32, cookie[24..28], 0xffff_ffff, .big);
    std.mem.writeInt(u32, cookie[28..32], 0, .big);

    var out: [voucher_size]u8 = @splat(0);
    std.mem.writeInt(u32, out[0..4], verification, .big);
    std.mem.writeInt(u32, out[4..8], voucher_size, .big);
    std.mem.writeInt(u32, out[8..12], key_type, .big);
    out[12] = keyChecksum(key);
    pc1.encrypt(out[16..48], &cookie, key);
    return out;
}

// -- Tests -----------------------------------------------------------------
//
// As in `pc1.zig`, the expected values come from compiling libmobi's reference
// routines and dumping their output, so these pin the derivation to a known
// good implementation.

const testing = std.testing;

test "crc register matches the reference" {
    // Deliberately not the standard CRC-32 values; see `crcRegister`.
    try testing.expectEqual(@as(u32, 0xca6598d0), crcRegister("abc"));
    try testing.expectEqual(@as(u32, 0x00000000), crcRegister(""));
    try testing.expectEqual(@as(u32, 0x2dfd2d88), crcRegister("123456789"));
}

test "crc register is not the standard crc-32" {
    // The standard CRC-32/ISO-HDLC check value for "123456789" is 0xCBF43926.
    // Getting these two confused silently produces wrong PIDs, so assert the
    // difference rather than leaving it to a comment.
    try testing.expect(crcRegister("123456789") != 0xcbf43926);
}

test "pid derivation matches the reference" {
    try testing.expectEqualStrings("YFEAZXZ*KZ", &try pidFromSerial("B0011T2ZQA123456"));
    try testing.expectEqualStrings("YQ8EQXN*IP", &try pidFromSerial("9876543210ABCDEF"));
    try testing.expectEqualStrings(
        "VVZEVVZEFQ",
        &try pidFromSerial("0123456789012345678901234567890123456789"),
    );
}

test "derived pids pass their own checksum" {
    try testing.expect(pidIsValid(try pidFromSerial("B0011T2ZQA123456")));
    try testing.expect(pidIsValid(try pidFromSerial("0123456789012345678901234567890123456789")));

    var bad = try pidFromSerial("B0011T2ZQA123456");
    bad[0] = if (bad[0] == 'A') 'B' else 'A';
    try testing.expect(!pidIsValid(bad));
}

test "rejects serials of the wrong length" {
    try testing.expectError(error.InvalidSerial, pidFromSerial("too short"));
    try testing.expectError(error.InvalidSerial, pidFromSerial("B0011T2ZQA1234567"));
    try testing.expectError(error.InvalidSerial, pidFromSerial(""));
}

test "pid key matches the reference" {
    var buf: [16]u8 = undefined;
    const cases = .{
        .{ "B0011T2ZQA123456", "c328c50c52ca079956f8a0283e7a0dca" },
        .{ "9876543210ABCDEF", "c33fedd298ce5b0182215edd66da9409" },
        .{ "0123456789012345678901234567890123456789", "ccca3d3e82aef34b42514519920e86b3" },
    };
    inline for (cases) |case| {
        const key = pidKey(try pidFromSerial(case[0]));
        try testing.expectEqualSlices(u8, try std.fmt.hexToBytes(&buf, case[1]), &key);
    }
}

test "key checksum matches the reference" {
    try testing.expectEqual(@as(u8, 0x36), keyChecksum(internal_reader_key));
    try testing.expectEqual(@as(u8, 0x9d), keyChecksum(internal_reader_key_v1));
}

test "a voucher round trips with the internal key" {
    const book_key: pc1.Key = "0123456789abcdef".*;
    const bytes = sealVoucher(book_key, internal_reader_key, key_type_internal, 0xdeadbeef);
    const voucher = Voucher.parse(&bytes);

    try testing.expectEqual(@as(u32, 0xdeadbeef), voucher.verification);
    try testing.expectEqual(@as(u32, voucher_size), voucher.size);

    const opened = openVoucher(voucher, internal_reader_key, key_type_internal) orelse
        return error.TestVoucherDidNotOpen;
    try testing.expectEqualSlices(u8, &book_key, &opened.key);
    try testing.expectEqual(@as(?u32, null), opened.expiresAt());
}

test "a voucher round trips with a device key" {
    const book_key: pc1.Key = "fedcba9876543210".*;
    const pid = try pidFromSerial("B0011T2ZQA123456");
    const device_key = pidKey(pid);

    const bytes = sealVoucher(book_key, device_key, key_type_device, 0x12345678);
    const voucher = Voucher.parse(&bytes);

    const opened = openVoucher(voucher, device_key, key_type_device) orelse
        return error.TestVoucherDidNotOpen;
    try testing.expectEqualSlices(u8, &book_key, &opened.key);

    // The wrong key, and the right key with the wrong type, must both fail.
    try testing.expectEqual(@as(?Unlocked, null), openVoucher(voucher, internal_reader_key, key_type_internal));
    try testing.expectEqual(@as(?Unlocked, null), openVoucher(voucher, device_key, key_type_internal));

    // A different device must not open it either.
    const other = pidKey(try pidFromSerial("9876543210ABCDEF"));
    try testing.expectEqual(@as(?Unlocked, null), openVoucher(voucher, other, key_type_device));
}