#!/usr/bin/env python3
"""Offline compile-profile regression; never invokes PlatformIO or a device."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import shlex
import subprocess
import tempfile

ROOT = Path(__file__).resolve().parents[2]
HERE = Path(__file__).resolve().parent
GUARD = ROOT / 'cmake/wolf_crypto_policy.h'
VENDOR = ROOT / 'managed_components/wolfssl__wolfssl'
ENV = {**os.environ, 'CCACHE_DISABLE': '1'}
SMALL = ('CURVE25519_SMALL', 'ED25519_SMALL')
VALIDATION = ('WOLFSSL_VALIDATE_ECC_IMPORT', 'WOLFSSL_ECDHX_SHARED_NOT_ZERO')
POLICY = (*SMALL, *VALIDATION)
ECC_BACKENDS = ('NO_ECC_CHECK_PUBKEY_ORDER', 'WOLF_CRYPTO_CB_ONLY_ECC',
                'WOLFSSL_ATECC508A', 'WOLFSSL_ATECC608A', 'WOLFSSL_CRYPTOCELL',
                'WOLFSSL_SILABS_SE_ACCEL', 'WOLFSSL_SE050', 'WOLFSSL_STM32_PKA')


def run(args, cwd=ROOT, **kw):
    return subprocess.run(args, cwd=cwd, env=ENV, timeout=60,
                          capture_output=True, text=True, **kw)


def require(result):
    if result.returncode:
        raise RuntimeError(result.stderr)
    return result.stdout


def clean(entry):
    args = entry.get('arguments') or shlex.split(entry['command'])
    result = []
    skip = False
    for arg in args:
        if skip:
            skip = False
        elif arg in ('-o', '-MF', '-MT', '-MQ'):
            skip = True
        elif arg not in ('-c', '-MD', '-MMD', '-MP'):
            result.append(arg)
    return result


def matrix():
    with tempfile.TemporaryDirectory(prefix='wolf-policy-') as tmp:
        tmp = Path(tmp)
        settings = tmp / 'wolfssl/wolfcrypt/settings.h'
        settings.parent.mkdir(parents=True)
        settings.write_text('/* Resolved settings supplied by matrix. */\n')
        base = ['HAVE_CURVE25519', 'HAVE_ED25519', 'HAVE_ECC',
                        'HAVE_ECC_CHECK_KEY', *POLICY]
        cases = [('valid', base, True)]
        cases += [(f'missing {m}', [x for x in base if x != m], False) for m in base]
        cases += [(m, base + [m], False) for m in
                  ('WOLFSSL_CURVE25519_BLINDING', 'HAVE_CURVE448', 'HAVE_ED448', *ECC_BACKENDS)]
        for label, defines, good in cases:
            p = run(['cc', '-x', 'c', '-fsyntax-only', '-I' + str(tmp),
                     '-include', str(GUARD), *['-D' + x for x in defines], '-'], input='')
            if (p.returncode == 0) != good or (not good and 'wolf crypto policy:' not in p.stderr):
                raise RuntimeError(f'guard matrix failed: {label}: {p.stderr}')
        print(f'PASS: {len(cases)} fail-closed guard cases')


def cmake_propagation():
    """Use the real policy module with a tiny stand-in IDF target graph."""
    with tempfile.TemporaryDirectory(prefix='wolf-cmake-') as tmp:
        tmp = Path(tmp)
        settings = tmp / 'wolfssl/wolfcrypt/settings.h'
        settings.parent.mkdir(parents=True)
        settings.write_text('\n'.join('#define ' + m for m in
                            ('HAVE_ECC', 'HAVE_ECC_CHECK_KEY', 'HAVE_CURVE25519',
                             'HAVE_ED25519', *SMALL)) + '\n')
        (tmp / 'unit.c').write_text('int main(void) { return 0; }\n')
        (tmp / 'CMakeLists.txt').write_text(f'''
cmake_minimum_required(VERSION 3.16)
project(wolf_policy_propagation C)
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
add_library(wolf STATIC unit.c)
target_include_directories(wolf PUBLIC "{tmp}")
function(idf_component_get_property out component prop)
    set(${{out}} wolf PARENT_SCOPE)
endfunction()
include("{ROOT / 'cmake/wolf_crypto_policy.cmake'}")
add_library(ssh STATIC unit.c)
target_link_libraries(ssh PUBLIC wolf)
add_executable(app unit.c)
target_link_libraries(app PRIVATE ssh)
''')
        require(run(['cmake', '-S', str(tmp), '-B', str(tmp / 'build')]))
        require(run(['cmake', '--build', str(tmp / 'build')]))
        entries = json.loads((tmp / 'build/compile_commands.json').read_text())
        if len(entries) != 3:
            raise RuntimeError('CMake policy fixture must compile library, SSH, app')
        for entry in entries:
            args = clean(entry)
            if not all('-D' + flag in args for flag in VALIDATION):
                raise RuntimeError('validation definitions failed PUBLIC propagation')
            if not any('wolf_crypto_policy.h' in arg for arg in args):
                raise RuntimeError('guard failed PUBLIC propagation')
        print('PASS: real CMake module PUBLIC definitions/guard across three targets')


def reviewed_ssh_body(name, body):
    """Independent allowlist of the reviewed parser delta, not generator output.

    Start from the hash-pinned original and require exact occurrence counts;
    the caller then compares the entire resulting function to the compiled file.
    """
    def replace(old, new, count=1):
        nonlocal body
        if body.count(old) != count:
            raise RuntimeError(f're-audit original parser anchor: {name}')
        body = body.replace(old, new)

    if name in ('DoUserAuthRequestEcc', 'DoUserAuthRequestEd25519'):
        replace('if (publicKeyTypeSz != pk->publicKeyTypeSz &&\n',
                'if (publicKeyTypeSz != pk->publicKeyTypeSz ||\n',
                2 if name == 'DoUserAuthRequestEcc' else 1)
    if name == 'DoUserAuthRequestEcc':
        replace('''    if (ret == WS_SUCCESS) {
        ret = GetStringRef(&rSz, &r, pk->signature, pk->signatureSz, &i);
    }

    if (ret == WS_SUCCESS) {
        ret = GetStringRef(&sSz, &s, pk->signature, pk->signatureSz, &i);
    }
''', '''    if (ret == WS_SUCCESS) {
        /* GetSize bounded sz by signatureSz - i: this end cannot wrap. */
        sz += i;
        ret = GetStringRef(&rSz, &r, pk->signature, sz, &i);
    }

    if (ret == WS_SUCCESS) {
        ret = GetStringRef(&sSz, &s, pk->signature, sz, &i);
    }

    if (ret == WS_SUCCESS && (i != sz || sz != pk->signatureSz))
        ret = WS_BUFFER_E;
''')
    elif name == 'DoUserAuthRequestEd25519':
        replace('''if (publicKeyTypeSz != pk->publicKeyTypeSz
                && WMEMCMP(publicKeyType,''',
                '''if (publicKeyTypeSz != pk->publicKeyTypeSz
                || WMEMCMP(publicKeyType,''')
        replace('''    if (ret == WS_SUCCESS) {
        ret = wc_ed25519_verify_msg_init(pk->signature + i, sz,''',
                '''    /* The signature string must consume the enclosing signature field. */
    if (ret == WS_SUCCESS && sz != pk->signatureSz - i)
        ret = WS_BUFFER_E;

    if (ret == WS_SUCCESS) {
        ret = wc_ed25519_verify_msg_init(pk->signature + i, sz,''')
    return body


def source_contract(database):
    pins = {
        ROOT / 'managed_components/wolfssl__wolfssh/src/internal.c':
            '81ff1f9166708abd5c2911e9fe57c0aee01c88b5d3f68c909ee8a856d37f36a9',
        VENDOR / 'wolfcrypt/src/ecc.c':
            '909c57e2756a8002df9f1d214483c659cb20db4a5f51047eda62d64ac458db06',
        VENDOR / 'wolfcrypt/src/curve25519.c':
            '9a0f6f0205245a8d19500a936d9b02bb71c8656713648408d1cb408362694b76',
        VENDOR / 'wolfcrypt/src/signature.c':
            '62ab3db3dfd251b2a2c73b69ef05aab6085d2e0d673fd9159514b3ee261cea4f',
    }
    for path, expected in pins.items():
        if hashlib.sha256(path.read_bytes()).hexdigest() != expected:
            raise RuntimeError(f're-audit key-validation source: {path}')
    entries = json.loads(database.read_text())
    entries = [e for e in entries if e['file'].endswith(
        '/security_overrides/wolfssh_internal/internal.c')]
    if len(entries) != 1:
        raise RuntimeError('expected one generated wolfSSH compilation input')
    generated = Path(entries[0]['file'])
    if not generated.is_absolute():
        generated = Path(entries[0]['directory']) / generated
    original = next(iter(pins)).read_text()
    derived = generated.read_text()
    for name in ('HashForId', 'KeyAgreeEcdh_server', 'KeyAgreeCurve25519_server',
                 'SignHEcdsa', 'DoUserAuthRequestEcc', 'DoUserAuthRequestEd25519',
                 'DoUserAuthRequestPublicKey'):
        pattern = rf'^(?:static )?(?:int|enum wc_HashType) {name}\('
        def extract(text):
            match = re.search(pattern, text, re.M)
            if match is None:
                raise RuntimeError(f'missing audited function {name}')
            end = text.index('\n}\n', match.start()) + 3
            return text[match.start():end]
        expected = reviewed_ssh_body(name, extract(original))
        if expected != extract(derived):
            raise RuntimeError(f're-audit modified generated crypto path: {name}')
    print('PASS: pinned originals; seven exact SSH paths including reviewed ECC/Ed parser deltas')
    print('INFO: generated wolfSSH SHA256 ' + hashlib.sha256(generated.read_bytes()).hexdigest())


def profiles(database, candidate):
    entries = json.loads(database.read_text())
    suffixes = ('wolfcrypt/src/ecc.c', 'wolfcrypt/src/signature.c',
                'wolfcrypt/src/curve25519.c', 'wolfcrypt/src/ed25519.c',
                'wolfcrypt/src/fe_operations.c', 'wolfcrypt/src/ge_operations.c',
                'wolfcrypt/src/fe_low_mem.c', 'wolfcrypt/src/ge_low_mem.c',
                'security_overrides/wolfssh_ssh/ssh.c',
                'security_overrides/wolfssh_internal/internal.c',
                'src/ssh_transport.c', 'src/ssh_security.c')
    for suffix in suffixes:
        matches = [e for e in entries if e['file'].endswith('/' + suffix)]
        if len(matches) != 1:
            raise RuntimeError(f'expected one compile entry for {suffix}: {len(matches)}')
        entry = matches[0]
        command = clean(entry)
        if candidate:
            command += ['-D' + x for x in POLICY] + ['-include', str(GUARD)]
        elif not any('wolf_crypto_policy.h' in x for x in command):
            raise RuntimeError(f'{suffix}: missing production guard; parent must reconfigure/build')
        text = require(run(command + ['-E', '-dM'], cwd=entry['directory']))
        macros = dict(re.findall(r'^#define (\w+)(?: (.*))?$', text, re.M))
        for name in (*POLICY, 'HAVE_ECC', 'HAVE_ECC_CHECK_KEY',
                     'HAVE_CURVE25519', 'HAVE_ED25519',
                     'WC_RNG_SEED_CB', 'WOLFSSL_ED25519_STREAMING_VERIFY',
                     'NO_WOLFSSL_ESP32_CRYPT_AES', 'NO_WOLFSSL_ESP32_CRYPT_HASH'):
            if name not in macros:
                raise RuntimeError(f'{suffix}: missing resolved {name}')
        if not any(x in macros for x in ('__XTENSA__', '__xtensa__')):
            raise RuntimeError('expected actual Xtensa compiler')
        if suffix == 'wolfcrypt/src/ecc.c' and 'HAVE_ECC_CHECK_PUBKEY_ORDER' not in macros:
            raise RuntimeError('ECC import validator has no software point check')
        for name in ('WOLFSSL_CURVE25519_BLINDING', 'HAVE_CURVE448', 'HAVE_ED448',
                     *ECC_BACKENDS):
            if name in macros:
                raise RuntimeError(f'{suffix}: unexpected {name}')
        require(run(command + ['-fsyntax-only'], cwd=entry['directory']))
        print(f'PASS: {"CANDIDATE replay" if candidate else "production"} macros + syntax: {suffix}')
    # Exercise failures with the real installed settings, not only fake headers.
    entry = matches[0]
    base = [x for x in clean(entry) if x not in ['-D' + m for m in POLICY]
            and 'wolf_crypto_policy.h' not in x]
    for missing in POLICY:
        command = base + ['-D' + x for x in POLICY if x != missing]
        command += ['-U' + missing, '-include', str(GUARD), '-E']
        p = run(command, cwd=entry['directory'])
        if p.returncode == 0 or 'wolf crypto policy:' not in p.stderr:
            raise RuntimeError(f'real settings accepted missing {missing}: {p.stderr}')
    print('PASS: real target settings reject each missing policy flag')


def vectors():
    with tempfile.TemporaryDirectory(prefix='wolf-vectors-') as tmp:
        tmp = Path(tmp)
        (tmp / 'user_settings.h').write_text('''
#define WOLFCRYPT_ONLY
#define NO_ASN
#define NO_RSA
#define NO_DH
#define NO_DSA
#define NO_AES
#define NO_DES3
#define NO_RC4
#define NO_MD4
#define NO_MD5
#define NO_SHA

#define NO_HMAC
#define WC_NO_RNG
#define NO_FILESYSTEM
#define NO_WRITEV
#define NO_DEV_RANDOM
#define NO_MAIN_DRIVER
#define HAVE_ECC
#define WOLFSSL_VALIDATE_ECC_IMPORT
#define WOLFSSL_ECDHX_SHARED_NOT_ZERO
#define HAVE_CURVE25519
#define HAVE_ED25519
#define CURVE25519_SMALL
#define ED25519_SMALL
#define WOLFSSL_SHA512
#define WOLFSSL_ED25519_STREAMING_VERIFY
''')
        sources = ['curve25519.c', 'ed25519.c', 'fe_operations.c',
                   'ge_operations.c', 'fe_low_mem.c', 'ge_low_mem.c', 'sha512.c']
        require(run(['cc', '-std=c99', '-O2', '-DWOLFSSL_USER_SETTINGS',
                     '-I' + str(tmp), '-I' + str(VENDOR), '-include', str(GUARD),
                     '-ffunction-sections', '-fdata-sections',
                     str(HERE / 'vectors.c'),
                     *[str(VENDOR / 'wolfcrypt/src' / s) for s in sources],
                     '-Wl,--gc-sections', '-o', str(tmp / 'vectors')]))
        print(require(run([str(tmp / 'vectors')])).strip())
        settings = tmp / 'user_settings.h'
        settings.write_text(settings.read_text().replace('#define WC_NO_RNG', '')
                            .replace('#define NO_DEV_RANDOM', '')
                                                        .replace('#define NO_FILESYSTEM', '')
                                                                                    .replace('#define NO_ASN', '') + '''
                                                        #define WOLFSSL_ASN_TEMPLATE
                                                        #define NO_CERTS
                                                        #define NO_PWDBASED
                                                        #define NO_PKCS12
#define USE_FAST_MATH
#define TFM_NO_ASM
#define TFM_TIMING_RESISTANT
#include <strings.h>
#define WOLFSSL_SMALL_STACK
#define ECC_TIMING_RESISTANT
#define NO_ECC_SIGN
#define SINGLE_THREADED
''')
        sources = ['ecc.c', 'tfm.c', 'wolfmath.c', 'random.c', 'sha256.c',
                           'sha512.c', 'memory.c', 'asn.c', 'hash.c', 'coding.c', 'signature.c']
        require(run(['cc', '-std=c99', '-O2', '-DWOLFSSL_USER_SETTINGS',
                     '-I' + str(tmp), '-I' + str(VENDOR), '-include', str(GUARD),
                     '-ffunction-sections', '-fdata-sections',
                     str(HERE / 'ecc_vectors.c'),
                     *[str(VENDOR / 'wolfcrypt/src' / s) for s in sources],
                     '-Wl,--gc-sections', '-o', str(tmp / 'ecc_vectors')]))
        print(require(run([str(tmp / 'ecc_vectors')])).strip())


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--compile-commands', type=Path,
                        default=ROOT / '.pio/build/esp32-s3-devkitc-1-n16r8/compile_commands.json')
    parser.add_argument('--candidate', action='store_true',
                        help='inject proposed flags/guard into old commands; NOT production evidence')
    parser.add_argument('--host-only', action='store_true')
    args = parser.parse_args()
    matrix()
    cmake_propagation()
    vectors()
    if not args.host_only:
        source_contract(args.compile_commands)
        profiles(args.compile_commands, args.candidate)


if __name__ == '__main__':
    main()