genpat.py

#!/usr/bin/python3
import glob
import logging
import os
import subprocess
import tempfile
import yaml
from contextlib import redirect_stdout

import elf


def c_string(s):
    return '"' + s.replace('\\', '\\\\').replace('"', '\\"') + '"'


def c_bytes(s):
    return '"' + ''.join('\\%03o' % ch for ch in s) + '"'


def run_nasm(source, defines={}, fmt='bin'):
    with tempfile.TemporaryDirectory() as d:
        asm_name = os.path.join(d, 'code.asm')
        out_name = os.path.join(d, 'code.bin')
        with open(asm_name, 'wt') as asm_file:
            asm_file.write(source)
        subprocess.check_call([
            'nasm', '-Iasm/', '-f', fmt, '-o', out_name,
            *('-D%s=%s' % d for d in defines.items()),
            asm_name])
        with open(out_name, 'rb') as bin_file:
            return bin_file.read()


def compile_pattern(source):
    defines = {
        'BEGIN': 'db 0x01',
        'REPEAT': 'db 0x02',
        'REPEAT0': 'db 0x03',
        'OPTION': 'db 0x04',
        'END': 'db 0x05',
        'ALT': 'db 0x06',
        'ANYBYTE': 0xF0,
        'ANYWORD': 0xF0F0,
        'ANYDWORD': 0xF0F0F0F0,
    }
    anti_defines = {
        'BEGIN': 'db 0x00',
        'REPEAT': 'db 0x00',
        'REPEAT0': 'db 0x00',
        'OPTION': 'db 0x00',
        'END': 'db 0x00',
        'ALT': 'db 0x00',
        'ANYBYTE': 0xFF,
        'ANYWORD': 0xFFFF,
        'ANYDWORD': 0xFFFFFFFF,
    }
    code1 = run_nasm(source, defines)
    code2 = run_nasm(source, anti_defines)
    return '.'.join(
        '0x%02X' % a if a == b else
        '(zlen'  if a == 1 else
        'zlen)+' if a == 2 else
        'zlen)*' if a == 3 else
        'zlen)?' if a == 4 else
        'zlen)'  if a == 5 else
        'zlen|zlen' if a == 6 else
        'any'    if a == 0xF0 else
        'error'
        for a, b in zip(code1, code2))


def compile_replacement(source):
    source = 'extern PORT\nbits 16\n' + source
    elf_object = run_nasm(source, fmt='elf')
    return elf.parse(elf_object)


def process_pattern(filename, pat):
    if 'name' not in pat:
        logging.error('%s: missing name', filename)
        return
    name = pat['name']

    if 'ragel' not in pat:
        if 'find' not in pat:
            logging.error('%s: %s: missing pattern',
                          filename, name)
            return

        try:
            pat['ragel'] = compile_pattern(pat['find'])
        except:
            logging.error('%s: %s: failed to assemble pattern',
                          filename, name, exc_info=True)
            return

        if 'error' in pat['ragel']:
            logging.error('%s: %s: something wrong with pattern',
                          filename, name)
            return

    if 'replace' in pat:
        try:
            code, exports, relocations = compile_replacement(pat['replace'])
        except:
            logging.error('%s: %s: failed to assemble replacement',
                          filename, name, exc_info=True)
            return

        global patch_count
        pat['patch'] = {
            'idx': patch_count,
            'code': code,
            'exports': exports,
            'relocations': relocations,
        }
        patch_count += 1

    return True


patch_count = 0
patterns = [pattern
            for f in sorted(glob.glob('patterns/*.yaml'))
            for pattern in yaml.safe_load_all(open(f))
            if process_pattern(f, pattern)]


def collect_symbols(patterns):
    symbols = {'PORT': 1}
    for pat in patterns:
        patch = pat.get('patch')
        if patch:
            for rel in patch['relocations']:
                symbols.setdefault(rel.r_sym, len(symbols) + 1)
            for s in patch['exports']:
                symbols.setdefault(s, len(symbols) + 1)
    return symbols


symbols = collect_symbols(patterns)


with redirect_stdout(open('patterns.rl', 'wt')):
    print("""/* Generated by genpat.py */
    %%{
    machine AdLibScanner;
    include x86 "x86.rl";
    main := |*
    """)
    for pat in patterns:
        print("(")
        print("#", pat['name'])
        print(pat['ragel'].strip())
        print(") => {")
        if 'patch' in pat:
            print("  MATCH(%s)" % pat['patch']['idx'])
        if 'warn' in pat:
            print("  WARN(%s)" % c_string(pat['warn']))
        if 'fatal' in pat:
            print("  FATAL(%s)" % c_string(pat['fatal']))
        print("};")
        print()
    print("""
    any;
    *|;
    }%%
    """)

with redirect_stdout(open('patch.inc', 'wt')):
    print("/* Generated by genpat.py */")
    print()
    print("#define SYMBOL_COUNT", len(symbols))
    print()
    print("const struct patch patch_data[] = {")
    for pat in patterns:
        patch = pat.get('patch')
        if patch:
            print("{")
            print("%s," % c_string(pat['name']))
            print("%s," % c_bytes(patch['code']))
            print("%s," % len(patch['code']))
            print("{")
            if patch['relocations']:
                for rel in patch['relocations']:
                    print("{ %s, %s, %s }," % (
                        rel.r_offset, symbols[rel.r_sym], rel.r_type))
            else:
                print("{ 0 }")
            print("},")
            print("{")
            if patch['exports']:
                for sym, value in patch['exports'].items():
                    print("{ %s, %s }," % (symbols[sym], value))
            else:
                print("{ 0 }")
            print("}")
            print("},")
    print("};")