# Copyright (c) 2026 Tiago Pinto <tiago@donttrustverify.pt>
# SPDX-License-Identifier: MIT
#
# Independent producer-side implementation of
# draft-fassbender-scitt-time-anchor-03, written from the published
# Internet-Draft alone. See README.md and CONSTRUCTION.md.
"""ots_wire.py: reading and writing the OpenTimestamps (.ots) wire format.

Written from scratch for the independent interoperability exercise on
profile draft-fassbender-scitt-time-anchor-03. Tree model plus
serialiser (OTS-SERIALIZE / OTS-DESERIALIZE, profile Section 1.4).

Format source: [OTS] python-opentimestamps v0.4.5 (sdist SHA-256
56726ccde97fb67f336a7f237ce36808e5593c3089d68d900b1c83d0ebf9dcfa,
identified in Section 5.2.1 of the profile). File:line references into
that release:

  - magic + version + hash_op + digest .... core/timestamp.py:273-341
  - recursive Timestamp (0xff continues; 0x00 attestation; otherwise an
    op tag followed by a Timestamp over op(msg))
    ...................................... core/timestamp.py:101-183
  - canonical ordering on serialisation:
      attestations sorted, all but the last prefixed 0xff 0x00;
      with ops present the last attestation also takes 0xff 0x00;
      ops sorted, all but the last prefixed 0xff
      ...................................... core/timestamp.py:101-127
      attestation order: TAG across classes; uri for pending; height
      for bitcoin; (TAG, payload) for unknown ones
      ...... core/notary.py:53-64,123-127,203-208,266-271
      op order: TAG, then argument ......... core/op.py:82-89
  - unsigned LEB128 varuint .............. core/serialize.py:142-156,189-199
  - binary ops 0xf0 append / 0xf1 prepend, varbytes arg 1..4096
    ...................................... core/op.py:73-121 (MAX_RESULT_LENGTH op.py:40)
  - unary ops 0xf2 reverse / 0xf3 hexlify . core/op.py:124-160
  - crypto ops 0x02 sha1 / 0x03 ripemd160 / 0x08 sha256 / 0x67 keccak256
    ...................................... core/op.py:163-230
  - attestations: 8-byte TAG + varbytes(payload) .... core/notary.py:19-92
      pending 83dfe30d2ef90c8e, payload = varbytes(uri)   notary.py:163,213-225
      bitcoin 0588960d73d71901, payload = varuint(height) notary.py:255,292-303
      unknown ones preserved as tag + payload            notary.py:86-94

Replay semantics follow the profile (Section 1.4, OTS-REPLAY):
append/prepend/sha256 only; any other operation leaves the branch
unevaluated (None). The parser knows the full tag set only in order to
traverse the stream.

Inline comments below are in Portuguese, the implementer's working
language.
"""

import hashlib
import io

MAGIC = b'\x00OpenTimestamps\x00\x00Proof\x00\xbf\x89\xe2\xe8\x84\xe8\x92\x94'
MAJOR_VERSION = 1
MAX_RESULT_LENGTH = 4096   # op.py:40
MAX_PAYLOAD_SIZE = 8192    # notary.py:26
MAX_URI_LENGTH = 1000      # notary.py:165

TAG_ATTESTATION = 0x00
TAG_CONTINUE = 0xff

BINARY_OPS = {0xf0: 'append', 0xf1: 'prepend'}
UNARY_OPS = {0xf2: 'reverse', 0xf3: 'hexlify'}
CRYPT_OPS = {0x02: ('sha1', 20), 0x03: ('ripemd160', 20),
             0x08: ('sha256', 32), 0x67: ('keccak256', 32)}
TAG_BY_NAME = {'append': 0xf0, 'prepend': 0xf1, 'reverse': 0xf2,
               'hexlify': 0xf3, 'sha1': 0x02, 'ripemd160': 0x03,
               'sha256': 0x08, 'keccak256': 0x67}

ATT_PENDING = bytes.fromhex('83dfe30d2ef90c8e')
ATT_BITCOIN = bytes.fromhex('0588960d73d71901')
ATT_LITECOIN = bytes.fromhex('06869a0d73d71b45')


class WireError(Exception):
    pass


# ---------- primitivas de leitura/escrita ----------

class Reader:
    def __init__(self, data: bytes):
        self.fd = io.BytesIO(data)

    def read(self, n: int) -> bytes:
        b = self.fd.read(n)
        if len(b) != n:
            raise WireError(f"truncado: pedi {n} bytes, obtive {len(b)}")
        return b

    def read_varuint(self) -> int:
        value, shift = 0, 0
        while True:
            b = self.read(1)[0]
            value |= (b & 0x7f) << shift
            if not (b & 0x80):
                return value
            shift += 7

    def read_varbytes(self, max_len: int, min_len: int = 0) -> bytes:
        n = self.read_varuint()
        if n > max_len:
            raise WireError(f"varbytes excede máximo: {n} > {max_len}")
        if n < min_len:
            raise WireError(f"varbytes abaixo do mínimo: {n} < {min_len}")
        return self.read(n)

    def at_eof(self) -> bool:
        pos = self.fd.tell()
        b = self.fd.read(1)
        self.fd.seek(pos)
        return b == b''


def write_varuint(out: bytearray, value: int):
    # serialize.py:142-156
    if value == 0:
        out.append(0)
        return
    while value != 0:
        b = value & 0x7f
        if value > 0x7f:
            b |= 0x80
        out.append(b)
        if value <= 0x7f:
            break
        value >>= 7


def write_varbytes(out: bytearray, data: bytes):
    write_varuint(out, len(data))
    out.extend(data)


# ---------- atestações ----------

class Attestation:
    def __init__(self, tag: bytes, payload: bytes):
        self.tag = tag
        self.raw_payload = payload
        self.kind = 'unknown'
        self.height = None
        self.uri = None
        p = Reader(payload)
        if tag == ATT_BITCOIN:
            self.kind = 'bitcoin'
            self.height = p.read_varuint()
            if not p.at_eof():
                raise WireError("payload bitcoin com bytes a mais")  # notary.py:91
        elif tag == ATT_PENDING:
            self.kind = 'pending'
            self.uri = p.read_varbytes(MAX_URI_LENGTH).decode('utf-8')
            if not p.at_eof():
                raise WireError("payload pending com bytes a mais")
        elif tag == ATT_LITECOIN:
            self.kind = 'litecoin'
            self.height = p.read_varuint()

    @classmethod
    def pending(cls, uri: str):
        payload = bytearray()
        write_varbytes(payload, uri.encode('utf-8'))   # notary.py:213-214
        return cls(ATT_PENDING, bytes(payload))

    @classmethod
    def bitcoin(cls, height: int):
        payload = bytearray()
        write_varuint(payload, height)                 # notary.py:292-293
        return cls(ATT_BITCOIN, bytes(payload))

    def sort_key(self):
        # notary.py:53-64 (TAG entre classes) + subordem por classe
        if self.kind == 'pending':
            sub = self.uri.encode('utf-8')             # notary.py:203-205
        elif self.kind in ('bitcoin', 'litecoin'):
            sub = self.height                          # notary.py:266-268
        else:
            sub = self.raw_payload                     # notary.py:123-125
        return (self.tag, sub if isinstance(sub, bytes) else bytes([]) + sub.to_bytes(8, 'big'))

    def serialize(self, out: bytearray):
        out.extend(self.tag)                           # notary.py:32-38
        write_varbytes(out, self.raw_payload)

    def __repr__(self):
        if self.kind == 'bitcoin':
            return f"BitcoinAttestation(height={self.height})"
        if self.kind == 'pending':
            return f"PendingAttestation({self.uri})"
        return f"Attestation({self.tag.hex()}, {len(self.raw_payload)}B)"


# ---------- operações ----------

class Op:
    def __init__(self, name: str, arg: bytes | None = None):
        self.name = name
        self.arg = arg
        self.tag = TAG_BY_NAME[name]

    def apply(self, msg: bytes) -> bytes | None:
        """OTS-REPLAY (§1.4): só append/prepend/sha256; resto -> None."""
        if self.name == 'append':
            return msg + self.arg
        if self.name == 'prepend':
            return self.arg + msg
        if self.name == 'sha256':
            return hashlib.sha256(msg).digest()
        return None

    def sort_key(self):
        # op.py:82-89: TAG, depois conteúdo do tuplo (arg)
        return (self.tag, self.arg or b'')

    def serialize(self, out: bytearray):
        out.append(self.tag)                           # op.py:41-42
        if self.tag in BINARY_OPS:
            write_varbytes(out, self.arg)              # op.py:92-94

    def __repr__(self):
        return self.name if self.arg is None else f"{self.name} {self.arg.hex()}"


def _read_op(r: Reader, tag: int) -> Op:
    if tag in BINARY_OPS:
        arg = r.read_varbytes(MAX_RESULT_LENGTH, min_len=1)  # op.py:96-100
        return Op(BINARY_OPS[tag], arg)
    if tag in UNARY_OPS:
        return Op(UNARY_OPS[tag])
    if tag in CRYPT_OPS:
        return Op(CRYPT_OPS[tag][0])
    raise WireError(f"tag de operação desconhecida: 0x{tag:02x}")


# ---------- árvore (Timestamp) ----------

class Node:
    """Estrutura recursiva do Timestamp: atestações + filhos (op, Node)."""
    def __init__(self):
        self.attestations: list[Attestation] = []
        self.children: list[tuple[Op, 'Node']] = []

    def serialize(self, out: bytearray):
        # timestamp.py:101-127, incluindo a ordenação canónica
        if not self.attestations and not self.children:
            raise WireError("um Timestamp vazio não é serializável")
        atts = sorted(self.attestations, key=Attestation.sort_key)
        if len(atts) > 1:
            for att in atts[:-1]:
                out.extend(b'\xff\x00')
                att.serialize(out)
        if not self.children:
            out.append(0x00)
            atts[-1].serialize(out)
            return
        if atts:
            out.extend(b'\xff\x00')
            atts[-1].serialize(out)
        kids = sorted(self.children, key=lambda c: c[0].sort_key())
        for op, node in kids[:-1]:
            out.append(0xff)
            op.serialize(out)
            node.serialize(out)
        op, node = kids[-1]
        op.serialize(out)
        node.serialize(out)

    @classmethod
    def deserialize(cls, r: Reader, depth: int = 0) -> 'Node':
        # timestamp.py:130-183
        if depth > 256:
            raise WireError("limite de recursão excedido")  # timestamp.py:186
        self = cls()

        def element(tag_byte: int):
            if tag_byte == TAG_ATTESTATION:
                att_tag = r.read(8)                          # notary.py:68
                payload = r.read_varbytes(MAX_PAYLOAD_SIZE)  # notary.py:70
                self.attestations.append(Attestation(att_tag, payload))
            else:
                op = _read_op(r, tag_byte)
                self.children.append((op, cls.deserialize(r, depth + 1)))

        tag = r.read(1)[0]
        while tag == TAG_CONTINUE:
            element(r.read(1)[0])
            tag = r.read(1)[0]
        element(tag)
        return self

    def walk(self, prefix=None):
        """Itera (caminho_de_ops, Attestation) por todos os ramos."""
        prefix = prefix or []
        for att in self.attestations:
            yield (list(prefix), att)
        for op, node in self.children:
            prefix.append(op)
            yield from node.walk(prefix)
            prefix.pop()


# ---------- ficheiro (DetachedTimestampFile) ----------

class Proof:
    def __init__(self, hash_op: str, digest: bytes, root: Node):
        self.hash_op = hash_op   # F2
        self.msg = digest        # F1
        self.root = root

    @property
    def branches(self):
        return list(self.root.walk())

    def bitcoin_branches(self):
        return [(ops, a) for ops, a in self.branches if a.kind == 'bitcoin']

    def pending_branches(self):
        return [(ops, a) for ops, a in self.branches if a.kind == 'pending']

    def replay(self, ops, start: bytes) -> bytes | None:
        v = start
        for op in ops:
            v = op.apply(v)
            if v is None:
                return None
            if len(v) > MAX_RESULT_LENGTH:               # op.py:22
                return None
        return v

    def serialize(self) -> bytes:
        # timestamp.py:316-323
        out = bytearray()
        out.extend(MAGIC)
        write_varuint(out, MAJOR_VERSION)
        out.append(TAG_BY_NAME[self.hash_op])
        out.extend(self.msg)
        self.root.serialize(out)
        return bytes(out)


def parse(data: bytes) -> Proof:
    """OTS-DESERIALIZE: bytes .ots -> Proof, ou WireError."""
    r = Reader(data)
    if r.read(len(MAGIC)) != MAGIC:                       # timestamp.py:329
        raise WireError("magic inválido")
    if r.read_varuint() != MAJOR_VERSION:                 # timestamp.py:331-333
        raise WireError("versão major não suportada")
    tag = r.read(1)[0]
    if tag not in CRYPT_OPS:                              # timestamp.py:335
        raise WireError(f"hash_op do ficheiro desconhecida: 0x{tag:02x}")
    name, dlen = CRYPT_OPS[tag]
    digest = r.read(dlen)                                 # timestamp.py:336
    root = Node.deserialize(r)
    if not r.at_eof():                                    # timestamp.py:339
        raise WireError("lixo após o fim da estrutura")
    return Proof(name, digest, root)