extract_and_scan.py

#!/usr/bin/env python

import os
import sys
import subprocess as sp
import tempfile
import hashlib
script_dir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(script_dir)
from fast_check_for_ffast_math import check_file

# MAX_SIZE = 30 * 1024 * 1024 * 1024 # 30GB

def dirsize(d):
    size = sp.run(['du', '-csb', d],
            universal_newlines=True,
            stdout=sp.PIPE, stderr=sp.DEVNULL
    ).stdout.splitlines()[-1].split()[0]
    return int(size)

def extract_with_unzip(wheel, dest):
    os.makedirs(dest, exist_ok=True)
    zip_out = sp.run(
        ['unzip', wheel, '*.so*', '-d', dest],
        stdout=sp.PIPE, stderr=sp.PIPE, check=False
    )

def remove(d):
    sp.run(['rm', '-rf', d], check=False)

def hash_file(path):
    BUF_SIZE = 65536
    sha256 = hashlib.sha256()
    with open(path, 'rb') as f:
        while True:
            data = f.read(BUF_SIZE)
            if not data:
                break
            sha256.update(data)
        return sha256.hexdigest()

def main(args):
    JOB_NUM = int(args[1])
    LOG_DIR = '/fastdata/pypi_wheels/logs'
    HASH_LOG = os.path.join(LOG_DIR, f'hashes.{JOB_NUM}.txt')
    MATH_LOG = os.path.join(LOG_DIR, f'ffast_math.{JOB_NUM}.txt')
    EXTRACT_DIR = '/fastdata/pypi_wheels/extract'
    os.makedirs(EXTRACT_DIR, exist_ok=True)
    os.makedirs(LOG_DIR, exist_ok=True)
    with open(HASH_LOG, 'w') as hash_log, open(MATH_LOG, 'w') as math_log:
        for file in args[2:]:
            with tempfile.TemporaryDirectory(
                prefix=os.path.join(EXTRACT_DIR, 'extract_')) as tmpdir:

                per_wheel_dir = os.path.join(tmpdir, file)
                extract_with_unzip(file, per_wheel_dir)
                for root, dirs, files in os.walk(per_wheel_dir):
                    for f in files:
                        path = os.path.join(root, f)
                        hash_log.write(f'{file}\t{hash_file(path)}\t{path}\n')
                        check_file(path, math_log)
            remove(per_wheel_dir)

if __name__ == '__main__':
    main(sys.argv)

fast_check_for_ffast_math.py

import sys
import mmap
from elftools.elf.elffile import ELFFile, ELFError
import struct

set_fast_math_code = bytes.fromhex('0fae5c24fc814c24fc408000000fae5424fcc3')

def load_bytes_from_elf(bindata, elf, vaddr, size):
    try:
        paddr = next(iter(elf.address_offsets(vaddr)))
    except StopIteration:
        return None
    return bindata[paddr:paddr+size]

def check_for_ffast_math(elf, bindata, addr):
    # +4 because we may have an endbr64 instruction (f30f1efa) in the prologue
    code = load_bytes_from_elf(bindata, elf, addr, len(set_fast_math_code)+4)
    if code is None:
        return False
    return set_fast_math_code in code

def check_file(filename, out=sys.stdout):
    print(f"{filename} ", end="", flush=True, file=out)
    try:
        try:
            elf = ELFFile(open(filename,'rb'))
        except ELFError:
            print("is not an ELF file", file=out)
            return
        if (arch := elf.get_machine_arch()) != 'x64':
            print(f"is {arch} not x64", file=out)
            return
        bindata = mmap.mmap(elf.stream.fileno(), 0, access=mmap.ACCESS_READ)
        try:
            s = elf.get_section_by_name('.init_array')
        except ELFError:
            print("error parsing ELF file", file=out)
            return
        if s is not None and (ia_data := s.data()):
            constructors = struct.unpack(f'<{len(ia_data)//8}Q', ia_data)
        else:
            constructors = []
        i = len(constructors)
        print(f"({i} constructor{'s'[:i^1]}) ", end="", flush=True, file=out)
        for addr in constructors:
            if check_for_ffast_math(elf, bindata, addr):
                print(f"contains ffast-math constructor at {hex(addr)}", file=out)
                break
        else:
            print("is clean", file=out)
    except Exception as e:
        print(f"error: {e}", file=out)

if __name__ == "__main__":
    for filename in sys.argv[1:]:
        check_file(filename)