# SPDX-FileCopyrightText: © 2026 Joe T. Sylve, Ph.D. <joe.sylve@gmail.com>
#
# SPDX-License-Identifier: AGPL-3.0-or-later

"""
Decryption support for PSpice® encrypted model files.

PSpice files are plain-text SPICE netlists with selected ``.SUBCKT``
and ``.MODEL`` blocks replaced by hex-encoded ciphertext between
``$CDNENCSTART`` / ``$CDNENCFINISH`` marker pairs.  Six encryption
modes are supported (0-5), spanning DES (modes 0-2) and AES-256 ECB
(modes 3-5).

See ``SPECIFICATIONS/pspice.md`` for the full specification.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from spice_crypt.pspice.keys import derive_keys, load_user_keys, mode_from_marker

if TYPE_CHECKING:
    import os
    from collections.abc import Generator

_PAD_SENTINEL = b" $jbs$"

# Bytes 62-63 of a non-final (overflow) block, written by the encoder when a
# plaintext line is longer than the 62-byte payload and continues into the
# next block.
_OVERFLOW_MARKER = b"$+"


def _make_cipher(mode: int, short_key: bytes):
    """Instantiate and key the correct cipher engine for *mode*."""
    if mode <= 2:
        from spice_crypt.pspice.des import PSpiceDES

        cipher = PSpiceDES()
        cipher.set_key(short_key)
        return cipher

    # Modes 3-5: AES-256 ECB
    from spice_crypt._aes import AES256ECB

    # Short key zero-padded to 32 bytes
    aes_key = (short_key + b"\x00" * 32)[:32]
    return AES256ECB(aes_key)


def _decrypt_64_block(cipher, mode: int, data: bytes) -> bytes:
    """Decrypt a single 64-byte block using the appropriate engine."""
    if mode <= 2:
        return cipher.process_block(data, decrypt=True)

    # AES: 4 x 16-byte ECB blocks
    result = bytearray(64)
    for i in range(4):
        result[i * 16 : i * 16 + 16] = cipher.decrypt_block(data[i * 16 : i * 16 + 16])
    return bytes(result)


def _split_block(block: bytes) -> tuple[bytes, bool]:
    """Split a decrypted 64-byte block into ``(content, line_complete)``.

    A single plaintext line may span several blocks.  The encoder fills each
    block with up to 62 bytes of line content (bytes 0-61); when content
    remains it writes the ``$+`` overflow marker into bytes 62-63 and
    continues in the next block.  The *final* block of a line is padded with
    the `` $jbs$`` sentinel followed by random fill, unless the content (plus
    its terminating carriage return) happens to fill the payload exactly.

    Returns the line content carried by this block and whether the block
    terminates the line.  Continuation blocks return ``line_complete=False``
    and their content is concatenated by the caller.
    """
    payload = block[:62]

    # Intact padding sentinel within the payload: content ends before it.
    idx = payload.find(_PAD_SENTINEL)
    if idx >= 0:
        return payload[:idx].rstrip(b"\x00"), True

    # Carriage-return line terminator (CRLF source files).  Everything after
    # it is padding — a possibly-truncated sentinel plus random fill.  Bytes
    # 62-63 are searched too: when the line content plus its ``\r`` fills the
    # block exactly (63 or 64 bytes), there is no room for the sentinel and the
    # terminator lands in the bytes normally used for the ``$+`` overflow
    # marker.  An overflow (continuation) block instead holds ``$+`` there, so
    # it has no ``\r`` and falls through to the overflow case below.
    cr = block.find(b"\r")
    if cr >= 0:
        return block[: cr + 1], True

    # Sentinel truncated at the payload/byte-62 boundary on a line with no
    # carriage return: only a prefix `` $jbs`` of the 6-byte sentinel survives
    # at the tail of the payload (bytes 62-63 are overwritten).  Confirm it is
    # genuine padding via the next sentinel byte, and require that this is not
    # an overflow block (whose bytes 62-63 are the ``$+`` marker).
    if block[62:64] != _OVERFLOW_MARKER:
        for n in range(len(_PAD_SENTINEL) - 1, 1, -1):
            if payload.endswith(_PAD_SENTINEL[:n]) and block[62] == _PAD_SENTINEL[n]:
                return payload[:-n].rstrip(b"\x00"), True

    # Overflow block: 62 bytes of content with more to come in the next block.
    return payload.rstrip(b"\x00"), False


class PSpiceFileParser:
    """Parser for PSpice encrypted files with streaming support.

    Reads a text-mode file object, passes non-encrypted lines through
    unchanged, and decrypts ``$CDNENCSTART`` / ``$CDNENCFINISH`` blocks.
    """

    def __init__(
        self,
        file_obj,
        user_key_file: str | os.PathLike | None = None,
        encrypted_file_path: str | os.PathLike | None = None,
        user_key_bytes: bytes | None = None,
    ):
        """
        Args:
            file_obj: Seekable text-mode file-like object.
            user_key_file: Path to user key CSV file for mode 4.
            encrypted_file_path: Path of the file being decrypted
                (used for user key identity matching in mode 4).
            user_key_bytes: Raw user key bytes for mode 4 (alternative
                to loading from a CSV file via *user_key_file*).
        """
        self.file_obj = file_obj
        self.user_key_file = user_key_file
        self.encrypted_file_path = encrypted_file_path
        self.user_key_bytes = user_key_bytes

    def decrypt_stream(self) -> Generator[bytes, None, tuple[int, int]]:
        """Stream-decrypt the file, yielding plaintext chunks.

        Non-encrypted lines are yielded as-is (encoded to bytes).
        Encrypted blocks are decrypted and the plaintext content is
        yielded line by line.

        Returns:
            Generator yielding bytes chunks.
            The return value (via ``StopIteration``) is ``(0, 0)``
            (PSpice files have no verification checksums).
        """
        line_buffer = b""
        in_encrypted_block = False
        cipher = None
        mode = 0
        is_header = False

        for line in self.file_obj:
            stripped = (
                line.strip() if isinstance(line, str) else line.decode("utf-8", "replace").strip()
            )

            # Check for block start marker
            if stripped.startswith("$CDNENCSTART"):
                in_encrypted_block = True
                is_header = True
                mode, version_str = mode_from_marker(stripped)

                # Derive keys
                user_key = self.user_key_bytes
                if user_key is None and mode == 4 and self.user_key_file:
                    user_key = load_user_keys(self.user_key_file, self.encrypted_file_path)
                short_key, _ = derive_keys(mode, version_str, user_key)
                cipher = _make_cipher(mode, short_key)
                continue

            # Check for block end marker
            if stripped.startswith("$CDNENCFINISH"):
                # Flush any pending (incomplete) line
                if line_buffer:
                    yield line_buffer + b"\n"
                    line_buffer = b""
                in_encrypted_block = False
                cipher = None
                continue

            if not in_encrypted_block:
                # Pass through non-encrypted lines
                line_bytes = line.encode("utf-8") if isinstance(line, str) else line
                yield line_bytes
                continue

            # Inside an encrypted block — decode and decrypt the hex line
            hex_str = stripped
            if not hex_str:
                continue

            try:
                block_data = bytes.fromhex(hex_str)
            except ValueError:
                continue

            if len(block_data) != 64:
                continue

            plaintext_block = _decrypt_64_block(cipher, mode, block_data)

            # Skip the encrypted header (first block after marker)
            if is_header:
                is_header = False
                continue

            # A plaintext line may span several blocks; accumulate content
            # until a block terminates the line.
            content, line_complete = _split_block(plaintext_block)
            line_buffer += content
            if line_complete:
                yield line_buffer + b"\n"
                line_buffer = b""

        # Flush any trailing line
        if line_buffer:
            yield line_buffer + b"\n"

        return (0, 0)
