#!/usr/bin/env python3
"""Compare PostgreSQL's test_saslprep(bytea) against an RFC 4013 oracle.

Requirements: Python 3, pip install 'psycopg[binary]', and PostgreSQL with
CREATE EXTENSION test_saslprep;

Examples:
  python saslprep_diff.py --dsn 'dbname=postgres' --smoke
  python saslprep_diff.py --dsn 'dbname=postgres' --all
  python saslprep_diff.py --show

Tested independently against Python's Unicode 3.2.0 tables. Requires a
PostgreSQL server and test_saslprep extension for database comparison.
"""
import argparse
import stringprep as sp
import unicodedata as ud

UCD = ud.ucd_3_2_0

# SASLprep (RFC 4013) prohibits tables C.1.2, C.2.1, C.2.2 and C.3-C.9.
# PostgreSQL also always rejects A.1 (unassigned in Unicode 3.2), so this
# oracle models PostgreSQL's stored-password policy, not a 'query' profile.
PROHIBITED = (
    sp.in_table_c12,
    sp.in_table_c21,
    sp.in_table_c22,
    sp.in_table_c3,
    sp.in_table_c4,
    sp.in_table_c5,
    sp.in_table_c6,
    sp.in_table_c7,
    sp.in_table_c8,
    sp.in_table_c9,
    sp.in_table_a1,
)


def saslprep_reference(s: str):
    """Return (status, expected bytea) for PostgreSQL's stored passwords.

    Where the B.1 and C.1.2 mapping tables overlap (U+200B), this oracle
    deliberately uses PostgreSQL's mapping priority: C.1.2 before B.1.
    That overlap deserves a separate cross-implementation test.
    """
    mapped = "".join(
        " " if sp.in_table_c12(ch) else "" if sp.in_table_b1(ch) else ch
        for ch in s
    )
    normalized = UCD.normalize("NFKC", mapped)

    if not normalized or any(
        table(ch) for ch in normalized for table in PROHIBITED
    ):
        return "PROHIBITED", None

    if any(sp.in_table_d1(ch) for ch in normalized):
        if (
            not sp.in_table_d1(normalized[0])
            or not sp.in_table_d1(normalized[-1])
            or any(sp.in_table_d2(ch) for ch in normalized)
        ):
            return "PROHIBITED", None

    return "SUCCESS", normalized.encode("utf-8")


def display(s):
    return " ".join(f"U+{ord(ch):04X}" for ch in s)


def candidates(all_scalars=False):
    # Focus on codepoints which normalize to a different bidi/category,
    # create trailing combining marks, or disappear during mapping.
    focus = [
        "\u0340", "\u0341", "\uFB1D", "\uFB2A", "\uFC5B",
        "\u2135", "\u2136", "\u2137", "\u2138", "\u20A8",
        "\u2103", "\u00AD", "\u00A0", "\u200B", "\u2168",
    ]
    contexts = [
        ("", ""), ("A", ""), ("", "A"),
        ("\u05D0", "\u05D1"), ("\u0627", "\u0628"),
        ("\u05D0", "1"), ("1", "\u05D1"),
        ("\u05D0", ""), ("", "\u05D1"),
    ]
    for middle in focus:
        for prefix, suffix in contexts:
            s = prefix + middle + suffix
            yield display(s), s

    # Baselines from RFC 4013 section 3 and bidi positive/negative controls.
    baselines = [
        "user", "USER", "I\u00ADX", "\u00AA", "\u2168", "\u0007",
        "\u06271", "\u05D01\u05D1", "\u05D0A\u05D1", "\u05D0\u20A8\u05D1",
        "\u05D0\u2135\u05D1", "A\u2135",
    ]
    for s in baselines:
        yield display(s), s

    if all_scalars:
        # All UTF-8-encodable Unicode scalar values, excluding U+0000.
        # The C API takes NUL-terminated strings, so a bytea containing
        # U+0000 is truncated by test_saslprep() and cannot be meaningfully
        # compared as a Unicode character with this harness.
        for cp in range(1, 0x110000):
            if 0xD800 <= cp <= 0xDFFF:
                continue  # Surrogates are not valid UTF-8 scalars.
            s = chr(cp)
            yield f"U+{cp:04X}", s


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dsn", default="dbname=postgres", help="psycopg DSN")
    mode = parser.add_mutually_exclusive_group()
    mode.add_argument("--smoke", action="store_true", help="targeted cases (default)")
    mode.add_argument("--all", action="store_true", help="include every Unicode scalar")
    mode.add_argument("--show", action="store_true", help="only print expected samples")
    parser.add_argument("--limit", type=int, default=30, help="max mismatches to print")
    args = parser.parse_args()

    if args.show:
        for index, (name, s) in enumerate(candidates()):
            if index >= args.limit:
                break
            status, result = saslprep_reference(s)
            print(name, s.encode().hex(), status,
                  result.hex() if result is not None else "NULL")
        return

    try:
        import psycopg
    except ImportError as exc:
        raise SystemExit("Install psycopg: pip install 'psycopg[binary]'") from exc

    with psycopg.connect(args.dsn) as conn, conn.cursor() as cur:
        cur.execute("""
            CREATE TEMP TABLE saslprep_cases (
                seq integer, label text, src bytea,
                expected_status text, expected_result bytea
            ) ON COMMIT DROP
        """)
        count = 0
        with cur.copy("""COPY saslprep_cases
            (seq, label, src, expected_status, expected_result)
            FROM STDIN""") as copy:
            for count, (name, s) in enumerate(candidates(args.all), start=1):
                status, expected = saslprep_reference(s)
                copy.write_row((count, name, s.encode("utf-8"), status, expected))

        # Materialize differences, so the PG function runs once per case.
        cur.execute("""
            CREATE TEMP TABLE saslprep_mismatches ON COMMIT DROP AS
            SELECT c.seq, c.label, encode(c.src, 'hex') AS src,
                   actual.status AS got_status,
                   encode(actual.res, 'hex') AS got_result,
                   c.expected_status,
                   encode(c.expected_result, 'hex') AS expected_result
            FROM saslprep_cases AS c
            CROSS JOIN LATERAL test_saslprep(c.src) AS actual
            WHERE (actual.status, actual.res)
                IS DISTINCT FROM (c.expected_status, c.expected_result)
        """)
        cur.execute("SELECT count(*) FROM saslprep_mismatches")
        mismatches = cur.fetchone()[0]
        print(f"Checked {count:,} cases; mismatches: {mismatches:,}")
        cur.execute("""
            SELECT label, src, got_status, got_result,
                   expected_status, expected_result
            FROM saslprep_mismatches ORDER BY seq LIMIT %s
        """, (args.limit,))
        for label, src, got_status, got, expected_status, expected in cur:
            print(f"{label}: input={src}, got={got_status}/{got}, "
                  f"expected={expected_status}/{expected}")


if __name__ == "__main__":
    main()