#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-only
"""Offline SHA-256 state-continuation teaching tool; no network operations.

Adapted from this project's test_editions SHA-256 reference implementation.
Not a production cryptographic library. Inputs and outputs are raw bytes or
hexadecimal; original_total_len includes every byte before original padding.
"""

import argparse
import hashlib
import json
import struct

_SHA256_K = (
    0x428A2F98,
    0x71374491,
    0xB5C0FBCF,
    0xE9B5DBA5,
    0x3956C25B,
    0x59F111F1,
    0x923F82A4,
    0xAB1C5ED5,
    0xD807AA98,
    0x12835B01,
    0x243185BE,
    0x550C7DC3,
    0x72BE5D74,
    0x80DEB1FE,
    0x9BDC06A7,
    0xC19BF174,
    0xE49B69C1,
    0xEFBE4786,
    0x0FC19DC6,
    0x240CA1CC,
    0x2DE92C6F,
    0x4A7484AA,
    0x5CB0A9DC,
    0x76F988DA,
    0x983E5152,
    0xA831C66D,
    0xB00327C8,
    0xBF597FC7,
    0xC6E00BF3,
    0xD5A79147,
    0x06CA6351,
    0x14292967,
    0x27B70A85,
    0x2E1B2138,
    0x4D2C6DFC,
    0x53380D13,
    0x650A7354,
    0x766A0ABB,
    0x81C2C92E,
    0x92722C85,
    0xA2BFE8A1,
    0xA81A664B,
    0xC24B8B70,
    0xC76C51A3,
    0xD192E819,
    0xD6990624,
    0xF40E3585,
    0x106AA070,
    0x19A4C116,
    0x1E376C08,
    0x2748774C,
    0x34B0BCB5,
    0x391C0CB3,
    0x4ED8AA4A,
    0x5B9CCA4F,
    0x682E6FF3,
    0x748F82EE,
    0x78A5636F,
    0x84C87814,
    0x8CC70208,
    0x90BEFFFA,
    0xA4506CEB,
    0xBEF9A3F7,
    0xC67178F2,
)


def padding(message_length):
    """Return SHA-256 padding for the given message length in bytes."""
    padding = b"\x80"
    padding += b"\x00" * ((56 - (message_length + 1) % 64) % 64)
    return padding + struct.pack(">Q", message_length * 8)


def extend(previous_hash_hex, original_total_len, suffix):
    """Continue a digest state and return the new hexadecimal SHA-256 digest."""
    if original_total_len < 0 or original_total_len >= 2**61:
        raise ValueError("original length must be a nonnegative byte count below 2**61")
    if len(bytes.fromhex(previous_hash_hex)) != 32:
        raise ValueError("SHA-256 digest must describe exactly 32 bytes")

    def rotr(value, shift):
        return ((value >> shift) | (value << (32 - shift))) & 0xFFFFFFFF

    def compress(state, block):
        words = list(struct.unpack(">16I", block))
        for index in range(16, 64):
            s0 = (
                rotr(words[index - 15], 7)
                ^ rotr(words[index - 15], 18)
                ^ (words[index - 15] >> 3)
            )
            s1 = (
                rotr(words[index - 2], 17)
                ^ rotr(words[index - 2], 19)
                ^ (words[index - 2] >> 10)
            )
            words.append((words[index - 16] + s0 + words[index - 7] + s1) & 0xFFFFFFFF)
        a, b, c, d, e, f, g, hh = state
        for index in range(64):
            big_s1 = rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25)
            choose = (e & f) ^ ((e ^ 0xFFFFFFFF) & g)
            t1 = (hh + big_s1 + choose + _SHA256_K[index] + words[index]) & 0xFFFFFFFF
            big_s0 = rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22)
            majority = (a & b) ^ (a & c) ^ (b & c)
            t2 = (big_s0 + majority) & 0xFFFFFFFF
            hh, g, f = g, f, e
            e = (d + t1) & 0xFFFFFFFF
            d, c, b, a = c, b, a, (t1 + t2) & 0xFFFFFFFF
        return tuple(
            (old + new) & 0xFFFFFFFF
            for old, new in zip(state, (a, b, c, d, e, f, g, hh))
        )

    state = struct.unpack(">8I", bytes.fromhex(previous_hash_hex))
    original_padding = padding(original_total_len)
    new_total = original_total_len + len(original_padding) + len(suffix)
    suffix_padding = b"\x80"
    suffix_padding += b"\x00" * ((56 - (len(suffix) + 1) % 64) % 64)
    suffix_padding += struct.pack(">Q", new_total * 8)
    payload = suffix + suffix_padding
    for offset in range(0, len(payload), 64):
        state = compress(state, payload[offset : offset + 64])
    return struct.pack(">8I", *state).hex()


def self_test():
    # Compare continuation against independent hashlib at padding boundaries.
    """Check continuation against hashlib at message-padding boundaries."""
    for length in (0, 1, 18, 55, 56, 63, 64, 65, 119, 120, 128):
        original = bytes(range(length))
        for suffix in (b"", b"extra", b"z" * 70):
            old_digest = hashlib.sha256(original).hexdigest()
            expected = hashlib.sha256(original + padding(length) + suffix).hexdigest()
            assert extend(old_digest, length, suffix) == expected
    print("33 continuation checks passed against hashlib")


def main():
    """Parse command-line arguments and run the tool."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    parser.add_argument("--digest", help="previous SHA-256 digest, 64 hex digits")
    parser.add_argument("--original-length", type=int, help="original byte count")
    parser.add_argument("--suffix-hex", help="additional bytes as hexadecimal")
    args = parser.parse_args()
    if args.self_test:
        self_test()
        return
    if args.digest is None or args.original_length is None or args.suffix_hex is None:
        parser.error("provide --digest, --original-length, and --suffix-hex")
    try:
        suffix = bytes.fromhex(args.suffix_hex)
        result = extend(args.digest, args.original_length, suffix)
        glue = padding(args.original_length)
    except (ValueError, struct.error) as error:
        parser.error(str(error))
    print(
        json.dumps(
            {
                "padding_hex": glue.hex(),
                "continued_digest": result,
                "processed_before_suffix": args.original_length + len(glue),
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()