mango/tools/kexttool/select.py

import yaml
import os
import re
from kexttool.kext import KextSource


selection_file_path = '.config/extensions.yaml'
kext_makefile_path = '.config/extensions.mk'


def description_to_list(kext):
    return list(filter(None, kext.description().split('\n')))


def create_selection_file():
    if not os.path.isdir('.config'):
        os.mkdir('.config')

    all_kexts = KextSource.scan(os.path.join(os.getcwd(), 'kexts'))
    if len(all_kexts) == 0:
        print('No kernel extensions available.')
        return

    out_file = open(selection_file_path, 'w')
    selection_info = {}

    for k in all_kexts:
        k_info = {
            'description': description_to_list(k),
            'build': None
        }

        selection_info[k.id()] = k_info

    out_file.write(yaml.dump(selection_info, default_flow_style=False))

    print('A list of available kexts has been created at {}'.format(selection_file_path))
    print('For each kext that you want to build, set the "build" key to:')
    print('  * internal: build the kext into the kernel image, to be loaded during boot.')
    print('  * external: build the kext into a standalone executable image, to be loaded at runtime.')
    print('  * null (or any other value): don\'t build the kext.')
    return


def find_kext_by_id(all_kexts, ident):
    for k in all_kexts:
        if k.id() == ident:
            return k

    return None


def external_kext_src_variable_name(kext, lang):
    return 'KEXT_{}_{}_SRC'.format(kext.id(underscores=True), lang)


def external_kext_obj_variable_name(kext):
    return 'KEXT_{}_OBJ'.format(kext.id(underscores=True))


def write_external_kext_source_list(fp, kext, lang):
    sources = kext.sources_filepath(lang)

    fp.write('{} := \\\n'.format(external_kext_src_variable_name(kext, lang)))

    for i, src in enumerate(sources):
        fp.write('\t{}'.format(src))

        if i != len(sources) - 1:
            fp.write(' \\')
        fp.write('\n')


def write_external_kext_variables(fp, kexts):
    for kext in kexts:
        fp.write('{} := \\\n'.format(external_kext_obj_variable_name(kext)))
        sources = kext.sources_filepath()

        for i, src in enumerate(sources):
            obj = os.path.splitext(src)[0] + '.o'
            fp.write('\t$(BUILD_DIR)/{}'.format(obj))

            if i != len(sources) - 1:
                fp.write(' \\')
            fp.write('\n')
    

def c_define_from_kext_id(ident):
    out = ident.upper()
    out = re.sub('[^0-9a-zA-Z]+', '_', out)
    return 'KEXT_{}'.format(out)


def create_kext_makefile():
    all_kexts = KextSource.scan(os.path.join(os.getcwd(), 'kexts'))
    selection_info = {}

    try:
        selection_fp = open(selection_file_path, 'r')
        selection_info = yaml.safe_load(selection_fp)
    except Exception as e:
        print('E: cannot load extension selection data ({})'.format(e))
        return

    try:
        makefile_fp = open(kext_makefile_path, 'w')
    except Exception as e:
        print('E: cannot write extension makefile ({})'.format(e))
        return

    internal_source_files = []
    internal_kexts = []
    external_kexts = []

    print('the following kernel extensions will be built:')
    for kext_id, build_info in selection_info.items():
        try:
            kext_info = find_kext_by_id(all_kexts, kext_id)
            if kext_info == None:
                print('W: cannot find kernel extension with ID "{}"'.format(kext_id))
                continue

            if 'build' not in build_info:
                continue

            if build_info['build'] != 'internal' and build_info['build'] != 'external':
                continue

            if build_info['build'] == 'internal':
                internal_kexts.append(kext_info)
                internal_source_files += kext_info.sources_filepath()
            elif build_info['build'] == 'external':
                external_kexts.append(kext_info)
            else:
                continue

            print('  * {} ({})'.format(kext_info.id(), build_info['build']))
        except Exception as e:
            print(e)
            continue

    if len(internal_source_files) > 0:
        makefile_fp.write('INTERNAL_KEXT_OBJ := \\\n')

    for i, src in enumerate(internal_source_files):
        obj = os.path.splitext(src)[0] + '.o'
        makefile_fp.write('\t$(BUILD_DIR)/{}'.format(obj))

        if i != len(internal_source_files) - 1:
            makefile_fp.write(' \\')
        makefile_fp.write('\n')

    if len(internal_kexts) > 0:
        makefile_fp.write('INTERNAL_KEXT_DEFINES := \\\n')

    for i, kext_info in enumerate(internal_kexts):
        makefile_fp.write('\t-D{}=1'.format(c_define_from_kext_id(kext_info.id())))

        if i != len(internal_kexts) - 1:
            makefile_fp.write(' \\')
        makefile_fp.write('\n')

    if len(internal_kexts) > 0:
        makefile_fp.write('INTERNAL_KEXT_INCLUDES := \\\n')

    for i, kext_info in enumerate(internal_kexts):
        makefile_fp.write('\t-I{}'.format(os.path.join('$(ROOT_DIR)', kext_info.src_dirpath(), 'include')))

        if i != len(selection_info.items()) - 1:
            makefile_fp.write(' \\')
        makefile_fp.write('\n')

    write_external_kext_variables(makefile_fp, external_kexts)
        
    for kext in external_kexts:
        makefile_fp.write('\nkexts/{}: $({})\n'.format(kext.id(), external_kext_obj_variable_name(kext)))
            

def select_kexts():
    if os.path.isfile(selection_file_path):
        create_kext_makefile()
    else:
        create_selection_file()