derivation.py

from __future__ import annotations

import hashlib
from dataclasses import dataclass
from enum import IntEnum

from .constants import (
    CENTRAL_STATE_DOMAIN_RECORDS,
    HASH_PREFIX,
    NAME_PROGRAM_ID,
    RECORDS_PROGRAM_ID,
    REVERSE_LOOKUP_CLASS,
    ROOT_DOMAIN_ACCOUNT,
)
from .public_key import find_program_address, public_key_seed


class RecordVersion(IntEnum):
    V1 = 1
    V2 = 2


@dataclass(frozen=True)
class DomainKey:
    pubkey: str
    hashed: bytes
    is_sub: bool
    parent: str | None = None
    is_sub_record: bool = False


def get_hashed_name(name: str) -> bytes:
    return hashlib.sha256((HASH_PREFIX + name).encode("utf-8")).digest()


def get_name_account_key(
    hashed_name: bytes,
    name_class: str | bytes | None = None,
    name_parent: str | bytes | None = None,
) -> str:
    seeds = [
        hashed_name,
        public_key_seed(name_class),
        public_key_seed(name_parent),
    ]
    pubkey, _bump = find_program_address(seeds, NAME_PROGRAM_ID)
    return pubkey


def _derive(
    name: str,
    parent: str = ROOT_DOMAIN_ACCOUNT,
    class_key: str | None = None,
) -> tuple[str, bytes]:
    hashed = get_hashed_name(name)
    pubkey = get_name_account_key(hashed, class_key, parent)
    return pubkey, hashed


def _record_class(record: RecordVersion | int | None) -> str | None:
    if record is None:
        return None
    return CENTRAL_STATE_DOMAIN_RECORDS if RecordVersion(record) == RecordVersion.V2 else None


def get_domain_key(domain: str, record: RecordVersion | int | None = None) -> DomainKey:
    if domain.endswith(".sol"):
        domain = domain[:-4]

    parts = domain.split(".")
    record_version = None if record is None else RecordVersion(record)
    record_class = _record_class(record_version)

    if len(parts) == 2:
        prefix = chr(int(record_version) if record_version is not None else 0)
        subdomain_name = prefix + parts[0]
        parent_key, _parent_hash = _derive(parts[1])
        pubkey, hashed = _derive(subdomain_name, parent_key, record_class)
        return DomainKey(pubkey=pubkey, hashed=hashed, is_sub=True, parent=parent_key)

    if len(parts) == 3 and record_version is not None:
        parent_key, _parent_hash = _derive(parts[2])
        sub_key, _sub_hash = _derive("\0" + parts[1], parent_key)
        record_prefix = "\x02" if record_version == RecordVersion.V2 else "\x01"
        pubkey, hashed = _derive(record_prefix + parts[0], sub_key, record_class)
        return DomainKey(
            pubkey=pubkey,
            hashed=hashed,
            is_sub=True,
            parent=parent_key,
            is_sub_record=True,
        )

    if len(parts) >= 3:
        raise ValueError("The domain is malformed")

    pubkey, hashed = _derive(domain, ROOT_DOMAIN_ACCOUNT)
    return DomainKey(pubkey=pubkey, hashed=hashed, is_sub=False)


def get_reverse_key(domain: str, is_sub: bool = False) -> str:
    domain_key = get_domain_key(domain)
    hashed_reverse_lookup = get_hashed_name(domain_key.pubkey)
    return get_name_account_key(
        hashed_reverse_lookup,
        REVERSE_LOOKUP_CLASS,
        domain_key.parent if is_sub else None,
    )


def get_record_key(domain: str, record: str) -> str:
    return get_domain_key(f"{record}.{domain}", RecordVersion.V1).pubkey