#!/usr/bin/env python3
"""Read installed SDK sources; compile extracted patched functions with host doubles.

No dependency writes or network. --build-dir additionally verifies a real IDF
build's Ninja source registration; it does not run a firmware build.
"""
# SPDX-License-Identifier: GPL-3.0-only
from __future__ import annotations
import argparse
from dataclasses import replace
import hashlib
import importlib.util
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import tempfile

sys.dont_write_bytecode = True
ROOT = Path(__file__).resolve().parents[2]
HERE = Path(__file__).resolve().parent
SPEC = importlib.util.spec_from_file_location("security_overrides", ROOT / "tools/security_overrides.py")
sdk = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = sdk
SPEC.loader.exec_module(sdk)
TLS_ENTRY = next(e for e in sdk.ENTRIES if e.name == "esp_tls_mbedtls")


def source_path(entry, idf, project=ROOT):
    return {"idf": idf, "project": project}[entry.root] / entry.source


AUXILIARY = {
    "components/esp_http_server/src/httpd_main.c": "a16ef65069dda13889c67b922f25eb566573983d6c24f01c089a902d5fd26149",
    "components/esp_http_server/src/httpd_txrx.c": "7659ad52c32f29b9a08208dc8b22d023edf274047835ed58107d82a47ccce00e",
}
FEATURES = ["MBEDTLS_SSL_PROTO_TLS1_2", "MBEDTLS_SSL_SRV_C",
            "MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED", "MBEDTLS_ECDH_C",
            "MBEDTLS_ECDSA_C", "MBEDTLS_AES_C", "MBEDTLS_GCM_C",
            "MBEDTLS_SHA256_C", "MBEDTLS_SHA384_C"]


def run(command, *, ok=True, cwd=None):
    env = dict(os.environ, CCACHE_DISABLE="1", PYTHONDONTWRITEBYTECODE="1",
               TMPDIR=str(ROOT / ".pio"))
    result = subprocess.run([str(x) for x in command], cwd=cwd, env=env,
                            capture_output=True, text=True, timeout=60)
    if (result.returncode == 0) != ok:
        raise AssertionError(f"command: {command}\n{result.stdout}\n{result.stderr}")
    return result.stdout + result.stderr


def extract(text, name):
    matches = list(re.finditer(r"^[A-Za-z_][\w* \t]*\b" + re.escape(name) + r"\([^;]*?\)\s*\{", text, re.M))
    assert len(matches) == 1, (name, len(matches))
    start = matches[0].start()
    brace = matches[0].end() - 1
    depth = 0
    tokens = re.finditer(r'/\*.*?\*/|//[^\n]*|"(?:\\.|[^"\\])*"|\'(?:\\.|[^\'\\])*\'|[{}]', text[brace:], re.S)
    for token in tokens:
        if token.group() == "{": depth += 1
        elif token.group() == "}":
            depth -= 1
            if depth == 0: return text[start:brace + token.end()] + "\n"
    raise AssertionError(name)


def typedef(text, name):
    match = re.search(r"typedef struct " + name + r"(?:_t)? \{.*?\} " + name + r"_t;", text, re.S)
    assert match, name
    return match.group() + "\n"


def expect_error(function, phrase):
    try:
        function()
    except (sdk.OverrideError, OSError) as error:
        assert phrase in str(error), str(error)
    else:
        raise AssertionError("expected rejection: " + phrase)


def generator_tests(idf, work):
    binary = work / "generated"
    manifest = sdk.generate(idf, ROOT, binary)
    before = {p: (p.read_bytes(), p.stat().st_mtime_ns) for p in binary.rglob("*") if p.is_file()}
    assert sdk.generate(idf, ROOT, binary) == manifest
    assert before == {p: (p.read_bytes(), p.stat().st_mtime_ns) for p in before}
    for entry in sdk.ENTRIES:
        original = source_path(entry, idf).read_bytes()
        derived = (binary / "security_overrides" / entry.name / Path(entry.source).name).read_bytes()
        assert derived.startswith(original[:original.index(b"*/") + 2])
        assert derived != original
    expect_error(lambda: sdk.apply_edits("x", (sdk.Edit("missing", "z"),)), "got 0")
    expect_error(lambda: sdk.apply_edits("xx", (sdk.Edit("x", "z"),)), "got 2")
    expect_error(lambda: sdk.generate(idf, ROOT, binary, ()), "absent")
    expect_error(lambda: sdk.generate(idf, ROOT, binary, (sdk.ENTRIES[0],) * 2), "duplicate")
    expect_error(lambda: sdk.generate(idf, ROOT, binary,
                 (sdk.ENTRIES[0], replace(sdk.ENTRIES[0], name="alias"))), "ambiguous")
    expect_error(lambda: sdk.generate(idf, ROOT, idf / "forbidden"), "separate")
    fake = work / "sdk"
    version = Path("components/esp_common/include/esp_idf_version.h")
    for rel in [version] + [Path(e.source) for e in sdk.ENTRIES if e.root == "idf"]:
        target = fake / rel; target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(idf / rel, target)
    last = fake / TLS_ENTRY.source
    last.write_bytes(last.read_bytes() + b"\n/* changed dependency */\n")
    failed = work / "failed"
    expect_error(lambda: sdk.generate(fake, ROOT, failed), "SHA256 mismatch")
    assert not failed.exists(), "must validate all inputs before output"
    # A failed regeneration must not silently update even the first derived file.
    expect_error(lambda: sdk.generate(fake, ROOT, binary), "SHA256 mismatch")
    assert before == {p: (p.read_bytes(), p.stat().st_mtime_ns) for p in before}
    shutil.copyfile(idf / TLS_ENTRY.source, last)
    (fake / version).write_text((fake / version).read_text().replace("VERSION_PATCH   0", "VERSION_PATCH   1"))
    expect_error(lambda: sdk.generate(fake, ROOT, failed), "5.5.0")
    shutil.copyfile(idf / version, fake / version)
    last.unlink()
    expect_error(lambda: sdk.generate(fake, ROOT, failed), "No such file")
    escaped = work / "escaped_output"; escaped.mkdir()
    (escaped / "security_overrides").symlink_to(fake, target_is_directory=True)
    expect_error(lambda: sdk.generate(idf, ROOT, escaped), "output escapes")
    print("Generator exact hashes/version/absent/ambiguous/atomic-plan/idempotence/path safety PASS")
    return binary


def extracted_tests(idf, binary, work):
    texts = {e.name: (binary / "security_overrides" / e.name / Path(e.source).name).read_text() for e in sdk.ENTRIES}
    aux = {}
    for rel, expected in AUXILIARY.items():
        raw = (idf / rel).read_bytes()
        assert hashlib.sha256(raw).hexdigest() == expected, rel
        aux[Path(rel).name] = raw.decode()
    https = texts["https_server"]
    functions = typedef(https, "httpd_ssl_ctx") + typedef(https, "httpd_ssl_transport_ctx")
    functions += extract(aux["httpd_main.c"], "httpd_stop")
    for name in ("security_override_wipe", "httpd_ssl_close", "httpd_ssl_open",
                 "free_secure_context", "create_secure_context", "httpd_ssl_start", "httpd_ssl_stop"):
        functions += extract(https, name)
    source = (HERE / "https.c").read_text().replace("/* SDK_FUNCTIONS */", functions)
    compile_run("https", source, work)
    scratch = texts["httpd_parse"]
    functions = "".join(extract(aux["httpd_txrx.c"], name) for name in ("httpd_recv_pending", "httpd_unrecv"))
    functions += "".join(extract(scratch, name) for name in ("security_override_wipe", "security_override_resize_scratch", "read_block", "httpd_req_cleanup"))
    compile_run("scratch", (HERE / "scratch.c").read_text().replace("/* SDK_FUNCTIONS */", functions), work)
    tls = texts["esp_tls_mbedtls"]
    original = (idf / TLS_ENTRY.source).read_text()
    assert extract(tls, "set_client_config") == extract(original, "set_client_config")
    assert extract(tls, "esp_create_mbedtls_handle") == extract(original, "esp_create_mbedtls_handle")
    guards = tls[tls.index("/* The server profile"):tls.index('static const char *TAG = "esp-tls-mbedtls";')]
    functions = extract(tls, "set_server_config") + extract(tls, "set_client_config")
    source = (HERE / "tls.c").read_text().replace("/* SDK_FUNCTIONS */", functions)
    source = source.replace("/* SDK_PKI */", typedef(tls, "esp_tls_pki")).replace("/* TLS_GUARDS */", guards)
    defines = ["-D" + f for f in FEATURES]
    compile_run("tls", source, work, defines + ["-DMBEDTLS_SSL_RENEGOTIATION", "-DCONFIG_MBEDTLS_SSL_RENEGOTIATION"])
    compile_run("tls_no_renegotiation", source, work, defines)
    # Compile actual injected guards independently of the behavioral doubles.
    guard_file = work / "guards.c"; guard_file.write_text(guards)
    for feature in FEATURES:
        run(["cc", "-E", "-x", "c", *["-D" + f for f in FEATURES if f != feature], guard_file], ok=False)
    run(["cc", "-E", "-x", "c", *defines, "-DCONFIG_MBEDTLS_DYNAMIC_BUFFER", guard_file], ok=False)
    print("TLS feature guard matrix (each required feature + dynamic buffer rejection) PASS")


def compile_run(name, source, work, flags=()):
    c = work / (name + ".c"); exe = work / name
    c.write_text("/* Extracted SDK sections retain their upstream Apache-2.0 license. */\n" + source)
    run(["cc", "-std=gnu11", "-O2", "-Wall", "-Wextra", "-Werror", "-Wno-unused-parameter",
         "-Wno-unused-function", "-Wno-unused-variable", *flags, "-I", HERE, c, "-o", exe])
    print(run([exe]).strip())


def cmake_fixture_tests(idf, work):
    # Use real component inputs with mock IDF target discovery. No SDK compilation.
    fixture = work / "cmake_fixture"; fixture.mkdir()
    lines = ["cmake_minimum_required(VERSION 3.18)", "project(security_fixture C)",
             f'set(TEST_IDF "{idf}")',
             'function(idf_build_get_property out property)',
             '  set(${out} "${TEST_IDF}" PARENT_SCOPE)', 'endfunction()',
             'function(idf_component_get_property out component property)',
             '  set(${out} "test_${component}" PARENT_SCOPE)', 'endfunction()']
    for e in sdk.ENTRIES:
        if e.root == "project":
            copied = fixture / e.source
            copied.parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(source_path(e, idf), copied)
        lines += [f'add_library(test_{e.component} STATIC "{source_path(e, idf, fixture)}")']
    lines += ['if(TEST_MISSING)', f'set_property(TARGET test_{sdk.ENTRIES[0].component} PROPERTY SOURCES missing.c)', 'endif()',
              'if(TEST_AMBIGUOUS)', f'set_property(TARGET test_{sdk.ENTRIES[0].component} APPEND PROPERTY SOURCES "{source_path(sdk.ENTRIES[0], idf, fixture)}")', 'endif()',
              'if(TEST_TARGET_MISSING)', 'function(idf_component_get_property out component property)',
              'set(${out} nonexistent PARENT_SCOPE)', 'endfunction()', 'endif()']
    for e in sdk.ENTRIES:
        lines += [f'set_source_files_properties("{source_path(e, idf, fixture)}" PROPERTIES COMPILE_FLAGS "-DSOURCE_FLAG" COMPILE_DEFINITIONS "SOURCE_DEFINE" COMPILE_OPTIONS "-fno-common")']
    lines += [f'include("{ROOT / "cmake/security_overrides.cmake"}")']
    for e in sdk.ENTRIES:
        lines += [f'file(GENERATE OUTPUT "${{CMAKE_BINARY_DIR}}/{e.name}.sources" CONTENT "$<TARGET_PROPERTY:test_{e.component},SOURCES>")',
                  f'get_property(flags SOURCE "${{SAK_SECURITY_{e.name}_GENERATED}}" PROPERTY COMPILE_FLAGS)',
                  'if(NOT flags STREQUAL "-DSOURCE_FLAG")', 'message(FATAL_ERROR "lost compile flags")', 'endif()',
                  f'get_property(inc SOURCE "${{SAK_SECURITY_{e.name}_GENERATED}}" PROPERTY INCLUDE_DIRECTORIES)',
                  f'if(NOT inc MATCHES "{source_path(e, idf, fixture).parent}")', 'message(FATAL_ERROR "lost original quoted include directory")', 'endif()']
    (fixture / "CMakeLists.txt").write_text("\n".join(lines) + "\n")
    build = work / "cmake_good"
    run(["cmake", "-G", "Ninja", "-S", fixture, "-B", build])
    for e in sdk.ENTRIES:
        source = (build / (e.name + ".sources")).read_text()
        assert source == str(build / "security_overrides" / e.name / Path(e.source).name)
    ninja = (build / "build.ninja").read_text()
    for path in [ROOT / "tools/security_overrides.py", idf / "components/esp_common/include/esp_idf_version.h"] + [source_path(e, idf, fixture) for e in sdk.ENTRIES]:
        assert str(path) in next(line for line in ninja.splitlines() if ": RERUN_CMAKE" in line), path
    for flag, phrase in (("TEST_MISSING", "found 0"), ("TEST_AMBIGUOUS", "found 2"), ("TEST_TARGET_MISSING", "missing component target")):
        output = run(["cmake", "-G", "Ninja", "-S", fixture, "-B", work / flag, "-D" + flag + "=ON"], ok=False)
        assert phrase in output, output
    print("CMake actual include: exact target replacement/properties/reconfigure/fail-closed matrix PASS")


def extension_fixture_tests(idf, work):
    # Prove the extension API, relative SOURCES in a child directory, real quoted
    # includes, target/source flags, and automatic fail-closed reconfiguration.
    fixture = work / "extension"
    for directory in ("cmake", "tools", "component/src", "component/include"):
        (fixture / directory).mkdir(parents=True, exist_ok=True)
    shutil.copyfile(ROOT / "cmake/security_overrides.cmake", fixture / "cmake/security_overrides.cmake")
    c = fixture / "component/src/example.c"
    c.write_text('#include "local.h"\n#include "extra.h"\n'
                 '#if !defined(SOURCE_FLAG) || !defined(SOURCE_DEFINE) || !defined(SOURCE_OPTION) || !defined(TARGET_DEFINE)\n'
                 '#error "compile properties were lost"\n#endif\n'
                 'int example(void) { return LOCAL + EXTRA + 1; }\n')
    original = c.read_bytes(); digest = hashlib.sha256(original).hexdigest()
    (c.parent / "local.h").write_text("#define LOCAL 10\n")
    (fixture / "component/include/extra.h").write_text("#define EXTRA 20\n")
    (fixture / "component/CMakeLists.txt").write_text('add_library(test_extension STATIC src/example.c)\n'
        'target_compile_definitions(test_extension PRIVATE TARGET_DEFINE)\n'
        'set_source_files_properties(src/example.c PROPERTIES COMPILE_FLAGS "-DSOURCE_FLAG" '
        'COMPILE_OPTIONS "-DSOURCE_OPTION" COMPILE_DEFINITIONS "SOURCE_DEFINE" '
        'COMPILE_DEFINITIONS_DEBUG "CONFIG_DEFINE" INCLUDE_DIRECTORIES "${CMAKE_CURRENT_SOURCE_DIR}/include")\n')
    (fixture / "main.c").write_text('int example(void); int main(void) { return example() != 32; }\n')
    wrapper = ('import sys\nfrom pathlib import Path\nsys.dont_write_bytecode = True\n'
               f'sys.path.insert(0, {str(ROOT / "tools")!r})\nimport security_overrides as sdk\n'
               'import argparse\np=argparse.ArgumentParser()\n'
               '[p.add_argument(a, type=Path, required=True) for a in ("--idf-path", "--project-dir", "--binary-dir")]\n'
               'a=p.parse_args()\n'
               f'e=sdk.Entry("extension", "extension", "project", "component/src/example.c", {digest!r}, '
               '(sdk.Edit("LOCAL + EXTRA + 1", "LOCAL + EXTRA + 2"),))\n'
               'sdk.generate(a.idf_path, a.project_dir, a.binary_dir, (e,))\n')
    (fixture / "tools/security_overrides.py").write_text(wrapper)
    (fixture / "CMakeLists.txt").write_text('cmake_minimum_required(VERSION 3.18)\nproject(extension C)\n'
        f'set(TEST_IDF "{idf}")\n'
        'function(idf_build_get_property out property)\nset(${out} "${TEST_IDF}" PARENT_SCOPE)\nendfunction()\n'
        'function(idf_component_get_property out component property)\nset(${out} "test_${component}" PARENT_SCOPE)\nendfunction()\n'
        'add_subdirectory(component)\ninclude(cmake/security_overrides.cmake)\n'
        'get_property(config_def SOURCE "${SAK_SECURITY_extension_GENERATED}" DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/component" PROPERTY COMPILE_DEFINITIONS_DEBUG)\n'
        'if(NOT config_def STREQUAL "CONFIG_DEFINE")\nmessage(FATAL_ERROR "lost per-config source definitions")\nendif()\n'
        'add_executable(check main.c)\ntarget_link_libraries(check PRIVATE test_extension)\n')
    build = work / "extension_build"
    run(["cmake", "-G", "Ninja", "-S", fixture, "-B", build])
    run(["cmake", "--build", build]); run([build / "check"])
    assert c.read_bytes() == original
    generated = build / "security_overrides/extension/example.c"
    stamp = generated.stat().st_mtime_ns
    run(["cmake", "--build", build]); assert generated.stat().st_mtime_ns == stamp
    c.write_bytes(original + b"\n/* upstream changed */\n")
    output = run(["cmake", "--build", build], ok=False)
    assert "SHA256 mismatch" in output, output
    assert generated.stat().st_mtime_ns == stamp
    print("Extension mapping + child relative source/includes/flags real compile + automatic mismatch rejection PASS")


def build_registration(build, idf):
    ninja = (build / "build.ninja").read_text()
    compile_lines = [line for line in ninja.splitlines() if ": C_COMPILER" in line]
    for e in sdk.ENTRIES:
        generated = build / "security_overrides" / e.name / Path(e.source).name
        matches = [line for line in compile_lines if str(generated) in line]
        assert len(matches) == 1, (e.name, matches)
        assert not any(str(source_path(e, idf)) in line for line in compile_lines), e.name
        assert generated.read_bytes() == sdk.render_entry(e, {"idf": idf, "project": ROOT})[1]
    print("Real IDF Ninja registration: each generated source once, originals absent, bytes verified PASS")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--idf-path", type=Path, default=Path.home() / ".platformio/packages/framework-espidf")
    parser.add_argument("--build-dir", type=Path)
    args = parser.parse_args()
    idf = args.idf_path.resolve()
    sdk.verify_version(idf)
    (ROOT / ".pio").mkdir(exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="sdk-security-", dir=ROOT / ".pio") as tmp:
        work = Path(tmp)
        binary = generator_tests(idf, work)
        extracted_tests(idf, binary, work)
        cmake_fixture_tests(idf, work)
        extension_fixture_tests(idf, work)
    if args.build_dir: build_registration(args.build_dir.resolve(), idf)
    print("SDK security overrides: all requested host checks PASS (not live TLS/hardware)")


if __name__ == "__main__":
    main()