Mbed TLS / TF-PSA-Crypto FFDH public key export buffer overflow

#!/usr/bin/env python3
# CVE-2026-34875 checker for Mbed TLS / TF-PSA-Crypto
# Exit codes: 0=PATCHED, 1=VULNERABLE, 2=UNKNOWN

import os
import re
import sys
import subprocess
from pathlib import Path

ROOT = Path(sys.argv[1] if len(sys.argv) > 1 else "/")

MBED_VULN_MIN = (3, 5, 0)
MBED_PATCHED = (3, 6, 6)
TF_VULN_EXACT = (1, 0, 0)
TF_PATCHED = (1, 1, 0)

results = []
notes = []


def norm_root(p):
    p = str(p)
    if ROOT == Path("/"):
        return p
    return str(ROOT / p.lstrip("/"))


def parse_semver(text):
    m = re.search(r'(\d+)\.(\d+)\.(\d+)', text or '')
    if not m:
        return None
    return tuple(int(x) for x in m.groups())


def cmp_ver(a, b):
    return (a > b) - (a < b)


def assess_mbed(version_text, source):
    ver = parse_semver(version_text)
    if not ver:
        results.append(("UNKNOWN", f"Mbed TLS version unreadable from {source}: {version_text}"))
        return

    # Distro/backport ambiguity: if package version has a distro suffix and upstream part is old,
    # the package may contain a backport not visible from the plain semver.
    if re.search(r'(ubuntu|deb|el\d|amzn|alpine|sles|fc\d)', version_text, re.I):
        results.append(("UNKNOWN", f"Mbed TLS distro package version needs vendor advisory validation: {source}={version_text}"))
        return

    if cmp_ver(ver, MBED_VULN_MIN) >= 0 and cmp_ver(ver, MBED_PATCHED) < 0:
        results.append(("VULNERABLE", f"Mbed TLS vulnerable version from {source}: {version_text}"))
    elif cmp_ver(ver, MBED_PATCHED) >= 0:
        results.append(("PATCHED", f"Mbed TLS patched version from {source}: {version_text}"))
    else:
        results.append(("PATCHED", f"Mbed TLS version older than affected range from {source}: {version_text}"))


def assess_tf(version_text, source):
    ver = parse_semver(version_text)
    if not ver:
        results.append(("UNKNOWN", f"TF-PSA-Crypto version unreadable from {source}: {version_text}"))
        return

    if re.search(r'(ubuntu|deb|el\d|amzn|alpine|sles|fc\d)', version_text, re.I):
        results.append(("UNKNOWN", f"TF-PSA-Crypto distro package version needs vendor advisory validation: {source}={version_text}"))
        return

    if ver == TF_VULN_EXACT:
        results.append(("VULNERABLE", f"TF-PSA-Crypto vulnerable version from {source}: {version_text}"))
    elif cmp_ver(ver, TF_PATCHED) >= 0:
        results.append(("PATCHED", f"TF-PSA-Crypto patched version from {source}: {version_text}"))
    else:
        results.append(("UNKNOWN", f"TF-PSA-Crypto unusual version from {source}: {version_text}"))


def read_file(path):
    try:
        with open(path, 'r', encoding='utf-8', errors='ignore') as f:
            return f.read()
    except Exception:
        return None


def check_headers():
    header_candidates = [
        norm_root('/usr/include/mbedtls/version.h'),
        norm_root('/usr/local/include/mbedtls/version.h'),
        norm_root('/opt/include/mbedtls/version.h'),
        norm_root('/include/mbedtls/version.h'),
    ]
    tf_candidates = [
        norm_root('/usr/include/tf-psa-crypto/version.h'),
        norm_root('/usr/local/include/tf-psa-crypto/version.h'),
        norm_root('/opt/include/tf-psa-crypto/version.h'),
        norm_root('/include/tf-psa-crypto/version.h'),
    ]

    for p in header_candidates:
        if os.path.isfile(p):
            data = read_file(p)
            if data:
                m = re.search(r'MBEDTLS_VERSION_STRING\s+"([^"]+)"', data)
                if m:
                    assess_mbed(m.group(1), p)
                    return True
    for p in tf_candidates:
        if os.path.isfile(p):
            data = read_file(p)
            if data:
                m = re.search(r'TF_PSA_CRYPTO_VERSION_STRING\s+"([^"]+)"', data)
                if m:
                    assess_tf(m.group(1), p)
                    return True
    return False


def run_cmd(cmd):
    try:
        out = subprocess.check_output(cmd, stderr=subprocess.DEVNULL, text=True)
        return out.strip()
    except Exception:
        return None


def check_pkg_config():
    if ROOT != Path('/'):
        return False
    hit = False
    for name in ['mbedtls', 'tf-psa-crypto']:
        out = run_cmd(['pkg-config', '--modversion', name])
        if out:
            hit = True
            if name == 'mbedtls':
                assess_mbed(out, f'pkg-config:{name}')
            else:
                assess_tf(out, f'pkg-config:{name}')
    return hit


def check_dpkg_status():
    status_path = norm_root('/var/lib/dpkg/status')
    if not os.path.isfile(status_path):
        return False
    data = read_file(status_path)
    if not data:
        return False
    hit = False
    for block in data.split('\n\n'):
        pkg = re.search(r'^Package:\s*(.+)$', block, re.M)
        ver = re.search(r'^Version:\s*(.+)$', block, re.M)
        if not pkg or not ver:
            continue
        name = pkg.group(1).strip()
        version = ver.group(1).strip()
        if name in ('libmbedtls-dev', 'libmbedtls21', 'libmbedtls14', 'libmbedtls12', 'libmbedcrypto16', 'libmbedcrypto7', 'libmbedcrypto3', 'mbedtls'):
            hit = True
            assess_mbed(version, f'dpkg:{name}')
        elif name in ('tf-psa-crypto', 'libtf-psa-crypto-dev'):
            hit = True
            assess_tf(version, f'dpkg:{name}')
    return hit


def check_apk_installed():
    path = norm_root('/lib/apk/db/installed')
    if not os.path.isfile(path):
        return False
    data = read_file(path)
    if not data:
        return False
    hit = False
    name = None
    ver = None
    for line in data.splitlines() + ['']:
        if line.startswith('P:'):
            name = line[2:].strip()
        elif line.startswith('V:'):
            ver = line[2:].strip()
        elif line == '':
            if name and ver:
                if 'mbedtls' in name:
                    hit = True
                    assess_mbed(ver, f'apk:{name}')
                elif 'tf-psa-crypto' in name:
                    hit = True
                    assess_tf(ver, f'apk:{name}')
            name = None
            ver = None
    return hit


def check_rpm_db():
    if ROOT != Path('/'):
        return False
    out = run_cmd(['rpm', '-qa', '--qf', '%{NAME} %{VERSION}-%{RELEASE}\n'])
    if not out:
        return False
    hit = False
    for line in out.splitlines():
        parts = line.split(None, 1)
        if len(parts) != 2:
            continue
        name, version = parts
        if 'mbedtls' in name:
            hit = True
            assess_mbed(version, f'rpm:{name}')
        elif 'tf-psa-crypto' in name:
            hit = True
            assess_tf(version, f'rpm:{name}')
    return hit


def summarize():
    if not results:
        print('UNKNOWN')
        print('No Mbed TLS or TF-PSA-Crypto evidence found in common package DBs or headers.')
        sys.exit(2)

    for status, msg in results:
        notes.append(f'[{status}] {msg}')

    has_vuln = any(status == 'VULNERABLE' for status, _ in results)
    has_patched = any(status == 'PATCHED' for status, _ in results)
    has_unknown = any(status == 'UNKNOWN' for status, _ in results)

    if has_vuln:
        print('VULNERABLE')
        for n in notes:
            print(n)
        sys.exit(1)
    if has_patched and not has_unknown:
        print('PATCHED')
        for n in notes:
            print(n)
        sys.exit(0)
    print('UNKNOWN')
    for n in notes:
        print(n)
    sys.exit(2)


if __name__ == '__main__':
    found = False
    found |= check_pkg_config()
    found |= check_dpkg_status()
    found |= check_apk_installed()
    found |= check_rpm_db()
    found |= check_headers()
    summarize()