A collection of Ghidra scripts and tools to improve UEFI Binary Analysis and automation
# Copyright 2023 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from collections import namedtuple
from copy import deepcopy
import json
import os
import shutil
import subprocess
from options import loadOptions


decompressProgram = namedtuple("Decompress_Program", ['name', 'type', 'path', 'pythonCommand', 'destination'])


def decompressBinary(options: decompressProgram, binary: str):
    ret = False
    dest = os.path.join(options.destination, f'{options.name}_{os.path.basename(binary)}')
    if options.type == 'UEFIExtract':
        ret = decompressUEFIExtract(options.path, binary)
        if ret:
            ret = parseUEFIExtract(binary, dest)
    elif options.type == 'CHIPSEC':
        ret = decompressCHIPSEC(options.pythonCommand, options.path, binary)
        if ret:
            ret = parseCHIPSEC(binary, dest)
    return ret


def decompressUEFIExtract(exePath, binPath):
    command = [exePath, binPath, 'unpack']
    proc = subprocess.run(command)
    if proc.returncode != 0:
        return False
    return True


# Decompress a binary using UEFIExtract
# Parse the output and copy the efi and te files
# into a folder a separate folder
def parseUEFIExtract(binPath, destName, fileType: str = 'All'):
    reportFile = f'{binPath}.report.txt'
    dump = binPath + '.dump'
    prePE = 'Section_PE32_image_'
    preTE = 'Section_TE_image_'
    post = '_body.bin'
    destPath = destName
    if not os.path.lexists(destPath):
        os.mkdir(destPath)
    if not os.path.isdir(destPath):
        return False
    results = {}
    with open(reportFile, 'r') as report:
        line = report.readline()
        while line:
            msplit = line.replace('  ', '').split('|')
            if msplit[0].strip() == 'File':
                if (fileType == 'All' and msplit[1].strip() not in ['Volume image', 'Pad', 'Freeform', 'Raw']) or (fileType == 'SMM' and msplit[1].strip() == 'SMM module'):
                    if len(msplit) == 7:
                        fname = msplit[5][msplit[5].find(' ', 1):].strip() + '_' + msplit[6].strip().replace(' ', '_')
                    else:
                        fname = msplit[5][msplit[5].find(' ', 1):].strip()
                    if msplit[5][msplit[5].find(' ', 1):].strip() not in results:
                        results[msplit[5][msplit[5].find(' ', 1):].strip()] = {}
                    oldpath = os.path.join(dump, prePE + fname + post)
                    newpath = os.path.join(destPath, fname)
                    if os.path.isfile(newpath):
                        newpath += "_1"
                    if not os.path.isfile(oldpath):
                        oldpath = os.path.join(dump, preTE + fname + post)
                        if not os.path.isfile(oldpath):
                            oldpath = None
                    if oldpath:
                        results[msplit[5][msplit[5].find(' ', 1):].strip()][newpath] = oldpath
                        shutil.copyfile(oldpath, newpath)
            line = report.readline()
    filemapjson = json.dumps(results, indent=2)
    with open(os.path.join(destPath, 'Map.json'), 'w') as f:
        f.write(filemapjson)
    return True


def decompressCHIPSEC(pythonCmd, exePath, binPath):
    command = [pythonCmd, os.path.join(exePath, 'chipsec_util.py'), '-n', '--skip_config', '-nl', 'uefi', 'decode', binPath]
    proc = subprocess.run(command)
    if proc.returncode != 0:
        return False
    return True


def parseCHIPSEC(binPath, destName, fileType: str = 'All'):
    destPath = os.path.join(os.path.dirname(binPath), destName)
    jsonPath = binPath + '.UEFI.json'
    if not os.path.lexists(destPath):
        os.mkdir(destPath)
    if not os.path.isdir(destPath):
        return False
    with open(jsonPath, 'r') as f:
        contents = f.read()
    chipDict = json.loads(contents)
    filemap = recursiveChipsec(chipDict, destPath)
    filemapjson = json.dumps(filemap, indent=2)
    with open(os.path.join(destPath, 'Map.json'), 'w') as f:
        f.write(filemapjson)


def recursiveChipsec(listObj: list, dest: str):
    results = {}
    for module in listObj:
        if 'children' in module:
            res = recursiveChipsec(module['children'], dest)
            for key in res.keys():
                if key not in results.keys():
                    results[key] = res[key]
                elif res[key] == results[key]:
                    continue
                else:
                    for i in res[key].keys():
                        if i not in results[key].keys():
                            results[key][i] = res[key][i]
        if 'class' in module and module['class'] == 'EFI_SECTION':
            if 'Type' in module and module['Type'] in [16, 17, 18, 22]:
                if module['parentGuid'] not in results.keys():
                    results[module['parentGuid']] = {}
                fname = module['parentGuid']
                if module['ui_string']:
                    fname += f'_{module["ui_string"]}'
                elif module['Name']:
                    fname += f'_{module["Name"]}'
                newpath = os.path.join(dest, fname)
                if os.path.isfile(newpath):
                    newpath += "_1"
                results[module['parentGuid']][newpath] = module['file_path']
                shutil.copyfile(module['file_path'], newpath)
    return results


def compareMaps(file1, file2):
    res = {}
    with open(file1, 'r') as f:
        contents = f.read()
    jsondata1 = json.loads(contents)
    with open(file2, 'r') as f:
        contents = f.read()
    jsondata2 = json.loads(contents)
    res['mismatchedFiles'] = []
    diff = list(set(jsondata1.keys()).difference(set(jsondata2.keys())))
    diff2 = list(set(jsondata2.keys()).difference(set(jsondata1.keys())))
    res['uniqueGUIDS1'] = deepcopy(diff)
    res['uniqueGUIDS2'] = diff2
    diff += diff2
    for key in jsondata1.keys():
        if key not in diff:
            if len(jsondata1[key].items()) != len(jsondata2[key].items()):
                res['mismatchedFiles'].append(key)
    return res


if __name__ == '__main__':
    options = loadOptions()
    tools = []
    binaries = []
    for key in options['Decompress']['Program']:
        Type = options['Decompress']['Program'][key].get('Type')
        Path = options['Decompress']['Program'][key].get('Path')
        Destination = options['Decompress']['Destination']
        PythonPath = options['Decompress'].get('Python', None)
        tools.append(decompressProgram(key, Type, Path, PythonPath, Destination))
    for file in options['Decompress']['Binary']['Files']:
        binaries.append(file)
    # add binaries from folders
    for folder in options['Decompress']['Binary']['Folders']:
        for file in os.listdir(folder):
            binaries.append(os.path.join(folder, file))
    # loop through binaries and parse which programs
    for bin in binaries:
        for tool in tools:
            decompressBinary(tool, bin)

    # compare if selected
    if options['Decompress']['Compare']:
        print('Beginning comparison ')
        keys = options['Decompress']['Program'].keys()
        excludelist = []
        if len(keys) > 1:
            for decompress1 in list(keys):
                excludelist.append(decompress1)
                for decompress2 in list(keys)[1:]:
                    if decompress2 in excludelist:
                        continue
                    print(f'Comparing {decompress1} and {decompress2}')
                    for file in binaries:
                        bn = os.path.basename(file)
                        file1 = os.path.join(options['Decompress']['Destination'], f'{decompress1}_{bn}', 'Map.json')
                        file2 = os.path.join(options['Decompress']['Destination'], f'{decompress2}_{bn}', 'Map.json')
                        res = compareMaps(file1, file2)
                        print(f'{decompress1}, {decompress2}, {file}, :')
                        print(res)

Tools/decompressHelper/decompress.py
dcg-red-team/UEFI_Surveyor
# Copyright 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# @category UEFISurveyor.headless

from collections import namedtuple
import yaml


decompressOptions = namedtuple("Decompress_Options", ['binaryPath', 'extractProgram', 'extractPath', 'extractType', 'extractDestination', 'pythonCommand'])


def loadOptions() -> dict:
    try:
        with open('options.yaml', 'r') as f:
            data = f.read()
    except IOError:
        data = ''
    try:
        options = yaml.safe_load(data)
    except yaml.YAMLError:
        options = {}
    return options


def getDecompressOptions(options: dict) -> decompressOptions:
    binaryPath = options['Binary']['Path']
    extractProgram = options['Binary']['Program']
    extractPath = options['Binary']['ProgramPath']
    extractType = options['Binary']['FileType']
    extractDest = options['Binary']['Destination']
    extractPython = options['Binary'].get('Python', None)
    decOptions = decompressOptions(binaryPath, extractProgram, extractPath, extractType, extractDest, extractPython)
    return decOptions

Tools/decompressHelper/options.py
dcg-red-team/UEFI_Surveyor
# Copyright 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from fnmatch import fnmatch
import json
import os
import sys
from typing import Optional
from uuid import UUID
import yaml


def parseDECfile(filename: str) -> str:
    # Extract the GUID section from a given 'DEC' file
    guid = False
    linein = ""

    with open(filename, 'r') as f:
        for line in f:
            if guid:
                linein += line
            if line.count('[') == 1 and line.count(']') == 1 and (line.find('Guid') > 0 or line.find('Ppis') > 0 or line.find('Protocols') > 0):
                guid = True
            elif line.count('[') == 1 and line.count(']') == 1:
                guid = False
    return linein


def normalizeName(inName: str) -> str:
    # Normalize a GUID name to match what is used within EFISeek
    # Add an _ between names within GUIDS and remove the first letter
    # example gEdkiiDynamicTablesPkgTokenSpaceGuid becomes Edkii_Dynamic_Tables_Pkg_Token_Space_Guid
    outName = inName[0:2]
    for count, letter in enumerate(inName[2:]):
        if letter.isupper() and not inName[count + 1].isupper():
            outName += f"_{letter}"
        else:
            outName += letter
    return outName.upper()


def normalizeGuid(inGuid: str) -> str:
    # Format Guid into a repesentation that can be input to UUID
    alpha = inGuid.lower().replace('{', '').replace(',', ' ').replace('}', '').replace('0x', '').strip()
    outGuid = ''
    i = 0
    for num in alpha.split(' '):
        if len(num) == 0:
            continue
        if i == 0:
            pad = 8
        elif i in [1, 2]:
            pad = 4
        else:
            pad = 2
        outGuid += f'{num.zfill(pad)}'
        i += 1
    return outGuid[:32]


def createDECDict(guid_strings: str, keytype: str = "GUID") -> dict:
    # Guids stored in the following format
    # gEdkiiDynamicTablesPkgTokenSpaceGuid = { 0xab226e66, 0x31d8, 0x4613, { 0x87, 0x9d, 0xd2, 0xfa, 0xb6, 0x10, 0x26, 0x3c } }

    guidDict = {}
    # For each new line get the name and GUID
    for line in guid_strings.split('\n'):
        if len(line) < 20 or line.count('=') == 0:
            continue
        nraw, graw = line.split('=')
        name = normalizeName(nraw.strip()[1:])
        guid = normalizeGuid(graw)
        mguid = f'{UUID(hex=guid)}'
        if keytype.upper() == "NAME":
            key = name
            val = mguid
        else:
            key = mguid
            val = name
        if key in guidDict.keys():
            print(f'Collision found in DEC: {key}, {val}, {guidDict[key]}')
        else:
            guidDict[key] = val
    return guidDict


def parseGUIDFile(filename: str) -> dict:
    # parse a GUID file from either EFISeek or UEFI Surveyor
    # The file will have repeating pattern of [section] followed
    # by entries of name/guid pairs
    linein = ""
    section = {}
    name = ''

    with open(filename, 'r') as f:
        for line in f:
            if line.count('[') == 1 and line.count(']') == 1:
                if name != '':
                    section[name] = linein
                linein = ''
                name = line.strip().replace('[', '').replace(']', '')
            else:
                if line.find('=') > 0:
                    linein += line
        section[name] = linein
    return section


def parseUEFIToolcsv(filename: str) -> dict:
    section = {}
    with open(filename, 'r') as f:
        for line in f:
            if len(line.split(',')) == 2:
                guid, desc = line.split(',')
                # if UUID(guid) in section.keys():
                #     section[UUID(guid)].append(desc)
                # else:
                #     section[UUID(guid)] = [desc]
                section[f'{UUID(guid)}'] = normalizeName(desc.strip())
    return section


def createGUIDDict(guid_strings: str, keytype: str = 'GUID', intype: str = 'GUID') -> dict:
    # Sort via keytype VARIABLE_GUID or NAME
    # VARIABLE_GUID: {72234213-0fd7-48a1-a59f-b41bc107fbcd} = VARIABLE_GUID
    # NAME: VARIABLE_GUID = {72234213-0fd7-48a1-a59f-b41bc107fbcd}
    guidDict = {}
    # For each new line get the name and GUID
    for line in guid_strings.split('\n'):
        if len(line) < 20 or line.count('=') == 0:
            continue
        if intype.upper() == 'NAME':
            name, graw = line.split('=')
        else:
            graw, name = line.split('=')
        tmpguid = graw.lower().replace('{', '').replace('}', '').strip()
        mguid = f'{UUID(tmpguid)}'
        name = name.strip()
        if keytype.upper() == "NAME":
            key = name
            val = mguid
        else:
            key = mguid
            val = name
        if key in guidDict.keys():
            print(f'Collision found in GUID DB: {key}, {val}, {guidDict[key]}')
        else:
            guidDict[key] = val
    return guidDict


def hasKey(dict: dict, key: str) -> Optional[str]:
    # check if a key exists within a dictionary of dictionaries
    for k in dict.keys():
        if key in dict[k].keys():
            return dict[k][key]
    return None


def combineDicts(dict1: dict, dict2: dict, section: str, combine: bool = False) -> None:
    # dict1 is a dictionary of dictionaries sorted by section
    # dict2 is a standard dictionary
    # The function will ensure the section exists then check the
    # dict1 before adding an entry
    if section not in dict1.keys():
        dict1[section] = {}
    for key in dict2.keys():
        isMatch = hasKey(dict1, key)
        print('ismatch', isMatch)
        if combine:
            if isMatch is None:
                dict1[section][key] = [dict2[key]]
            elif dict2[key].strip() not in isMatch:
                isMatch.append(dict2[key].strip())
        else:
            if isMatch is None:
                dict1[section][key] = dict2[key]
            elif isMatch.strip() != dict2[key].strip():
                print(f'Collision found when combining dictionaries: {key}, {isMatch}, {dict2[key]}')


if __name__ == '__main__':
    # open options file and parse
    _file = open('options.yaml')
    data = _file.read()
    _file.close()
    options = yaml.safe_load(data)
    if options['GUIDS']['OUTPUT'] not in ['EFISeek', 'UEFISurveyor']:
        print('Missing output type')
        sys.exit()
    # parse input and combine into output format
    outDict = {}
    if 'COMBINED' in options['GUIDS']:
        combined = options['GUIDS']['COMBINED']
    else:
        combined = False
    print(combined, type(combined))
    for key in options['GUIDS']['INPUT'].keys():
        if options['GUIDS']['INPUT'][key]['TYPE'] == 'DEC':
            if not options['GUIDS']['INPUT'][key]['OPTIONS']:
                for dirpath, dirs, files in os.walk(options['GUIDS']['INPUT'][key]['PATH']):
                    for filename in files:
                        fname = os.path.join(dirpath, filename)
                        if fname.endswith('.dec'):
                            print(f'found file {fname}')
                            tmp = parseDECfile(fname)
                            tmpDB = createDECDict(tmp, 'GUID')
                            combineDicts(outDict, tmpDB, 'EDK', combined)
            else:
                for name in options['GUIDS']['INPUT'][key]['OPTIONS']:
                    filepath = os.path.join(options['GUIDS']['INPUT'][key]['PATH'], name)
                    dec_files = [f.name for f in sorted(os.scandir(filepath), key=lambda x: x.name)
                                 if fnmatch(f.name, '*.dec')]
                    print(f'dec_files {dec_files}')
                    for dec in dec_files:
                        tmp = parseDECfile(os.path.join(filepath, dec))
                        tmpDB = createDECDict(tmp, 'NAME' if options['GUIDS']['OUTPUT'] == 'EFISeek' else 'GUID')
                        combineDicts(outDict, tmpDB, key, combined)
        elif options['GUIDS']['INPUT'][key]['TYPE'] in ['EFISeek', 'UEFISurveyor']:
            sections = parseGUIDFile(options['GUIDS']['INPUT'][key]['PATH'])
            for section in sections.keys():
                tmpDB = createGUIDDict(sections[section], 'NAME' if options['GUIDS']['OUTPUT'] == 'EFISeek' else 'GUID', 'NAME' if options['GUIDS']['INPUT'][key]['TYPE'] == 'EFISeek' else 'GUID')
                combineDicts(outDict, tmpDB, section, combined)
        elif options['GUIDS']['INPUT'][key]['TYPE'] == 'UEFITool':
            sections = parseUEFIToolcsv(options['GUIDS']['INPUT'][key]['PATH'])
            combineDicts(outDict, sections, 'UEFITOOL', combined)

    # write the output into tmp_guid_db file
    with open("tmp_guid_db", 'w') as outfile:
        if options['GUIDS']['OUTPUT'] == 'EFISeek':
            for sect in outDict:
                outfile.writelines(f'[{sect}]\n')
                for key in outDict[sect].keys():
                    outfile.writelines(f'{key} = {{{outDict[sect][key]}}}\n')
        elif options['GUIDS']['OUTPUT'] == 'UEFISurveyor':
            newOut = {}
            for sect in outDict:
                newOut.update(outDict[sect])
            outfile.write(json.dumps(newOut, indent=4))
        elif options['GUIDS']['OUTPUT'] == 'UEFITool':
            for sect in outDict:
                for key in outDict[sect].keys():
                    outfile.writelines(f'{key},{outDict[sect][key]}\n')

    print('\n\nCreated tmp_guid_db')

Tools/guidFinder/guid_gen.py
dcg-red-team/UEFI_Surveyor
# Copyright 2023 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# This script identifies module entry point and labels
# efi specific functions, protocols, and guids
# @category UEFISurveyor.headless

from argparse import ArgumentParser
from fnmatch import fnmatch
import json
import os
import sys
from typing import Tuple
from copy import deepcopy


def getDictionary(filename: str) -> dict:
    with open(filename, 'r') as f:
        contents = f.read()
    return json.loads(contents)


def getDiff(seq1, seq2) -> Tuple[list, list]:
    diff = set(seq1).difference(set(seq2))
    diff2 = set(seq2).difference(set(seq1))
    return (list(diff), list(diff2))


def getDf(j1: dict, j2: dict) -> dict:
    res = {}
    (diff, diff2) = getDiff(j1.keys(), j2.keys())
    res['uniqueFunctions1'] = deepcopy(diff)
    res['uniqueFunctions2'] = diff2
    diff += diff2
    res['mismatchedHashes'] = []
    for key in j1.keys():
        if key not in diff:
            if j1[key] != j2[key]:
                res['mismatchedHashes'].append(key)
    return res


def isMismatched(DFDict):
    ret = False
    if DFDict['uniqueFunctions1'] or DFDict['uniqueFunctions2'] or DFDict['mismatchedHashes']:
        ret = True
    return ret


def compareDirs(dirpath1, dirpath2):
    res = {}
    json_files = [f.name for f in sorted(os.scandir(dirpath1), key=lambda x: x.name)
                  if fnmatch(f.name, '*json')]
    json_files2 = [f.name for f in sorted(os.scandir(dirpath2), key=lambda x: x.name)
                   if fnmatch(f.name, '*json')]
    diff, diff2 = getDiff(json_files, json_files2)
    res['uniqueFiles1'] = deepcopy(diff)
    res['uniqueFiles2'] = diff2
    res['misMatchedFiles'] = []
    diff += diff2
    for f in json_files:
        if f not in diff:
            j1 = os.path.join(dirpath1, f)
            j2 = os.path.join(dirpath2, f)
            res[f] = compareFiles(j1, j2)
            if isMismatched(res[f]):
                res['misMatchedFiles'].append(f)
    return res


def writeHashFile(results):
    jfile = 'hashlog.json'
    with open(jfile, 'w') as jf:
        jf.write(json.dumps(results, indent=4))


def compareFiles(filepath1, filepath2):
    jsondata1 = getDictionary(filepath1)
    jsondata2 = getDictionary(filepath2)
    diff = getDf(jsondata1, jsondata2)
    return diff


if __name__ == '__main__':
    parser = ArgumentParser(prog='comparefile')
    subparsers = parser.add_subparsers()

    # directory command args
    parser_dir = subparsers.add_parser('directory')
    parser_dir.add_argument('arg1', metavar='dir1', type=str, help='filepath of directory')
    parser_dir.add_argument('arg2', metavar='dir2', type=str, help='filepath of directory')
    parser_dir.set_defaults(func=compareDirs)

    # file command args
    parser_file = subparsers.add_parser('file')
    parser_file.add_argument('arg1', metavar='file1', type=str, help='filepath of file')
    parser_file.add_argument('arg2', metavar='file2', type=str, help='filepath of file')
    parser_file.set_defaults(func=compareFiles)

    if len(sys.argv) > 1:
        args = parser.parse_args(sys.argv[1:])
        print(args)
        res = args.func(args.arg1, args.arg2)
        writeHashFile(res)
    else:
        parser.parse_args(['help'])

Tools/hashCompare/compareHashes.py
dcg-red-team/UEFI_Surveyor
# Copyright 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
from typing import Union
import yaml


def getRequired(lpath: str, entry: Union[dict, list, str]) -> list:
    rl = []
    if isinstance(entry, dict):
        for key in entry.keys():
            rl += getRequired(os.path.join(lpath, key), entry[key])
    if isinstance(entry, list):
        for key in entry:
            rl += getRequired(lpath, key)
    if isinstance(entry, str):
        rl.append(os.path.join(lpath, entry))
    return rl


def genHeaders(pkg_list: list, localpath: str, nodirs: list) -> str:
    ret = ''
    for pkg in pkg_list:
        for root, dirs, files in os.walk(os.path.join(localpath, pkg, "Include")):
            for file in files:
                if file.endswith('.h') and not any(root.endswith(i) for i in nodirs):
                    ret += "{}\n".format(os.path.join(root, file))
    return ret


def genIncludes(pkg_list: list, localpath: str) -> str:
    ret = ''
    for pkg in pkg_list:
        ret += '{}\n'.format(os.path.join(localpath, pkg, "Include"))
    return ret


if __name__ == '__main__':
    ARCHS = ["X64", "Arm", "Ebc", "Ia32", "RiscV", "RiscV64", "AArch64"]
    # read in options.yaml
    _file = open('options.yaml')
    data = _file.read()
    options = yaml.safe_load(data)
    # ensure ARCH is within know architectures
    if 'ARCH' not in options or options['ARCH'] not in ARCHS:
        print(f"Missing valid architecture please provide one from the list {' '.join(i for i in ARCHS)}")
        sys.exit()
    # list of directories not to include
    nodirs = [i for i in ARCHS if i != options['ARCH']]
    # set name of outfile
    filename = 'uefi_{}.prf'.format(options['ARCH']).lower()
    # check to ensure configuration contains proper information
    if 'HEADERS' not in options or 'EDK2' not in options['HEADERS']:
        print('Missing EDK2')
        sys.exit()
    if 'PKGS' not in options['HEADERS']['EDK2']:
        print('Missing PKGS')
        sys.exit()
    # create prf file
    with open(filename, 'w') as outfile:
        for req in getRequired(options['HEADERS']['EDK2']['PATH'], options['REQUIRED']):
            outfile.writelines(f'{req}\n')
        for mname in options['HEADERS'].keys():
            outfile.writelines(genHeaders(options['HEADERS'][mname]['PKGS'], options['HEADERS'][mname]['PATH'], nodirs))
        outfile.writelines('\n\n')
        outfile.writelines('{}\n'.format(os.path.join(options['HEADERS']['EDK2']['PATH'], "MdePkg", "Include", options['ARCH'])))
        outfile.writelines('{}\n'.format(os.path.join(options['HEADERS']['EDK2']['PATH'], "MdePkg", "Include", "Library")))
        for mname in options['HEADERS'].keys():
            outfile.writelines(genIncludes(options['HEADERS'][mname]['PKGS'], options['HEADERS'][mname]['PATH']))
    print(f'Successfully generated {filename}')

Tools/prfGenerator/gen_prf.py
dcg-red-team/UEFI_Surveyor
# Copyright 2023 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# @category UEFISurveyor.internal

from copy import deepcopy
from artifacts import artifacts
from ghidra_funcs import varnodeConverter, GhidraUtils
from ghidra.program.model.pcode import PcodeOp

from guids import convertGuidStr, guids
from ghidra.program.model.mem import MemoryAccessException
from ghidra.program.model.address import AddressSpace
from ghidra.program.model.data import UnicodeDataType

from uuid import UUID
import struct


class EFIUtils(GhidraUtils):
    # UEFI Specific functions needed for binary analysis
    def __init__(self, currentProgram):
        super(EFIUtils, self).__init__(currentProgram)

    def labelModuleEntryPoints(self, eps):
        # update list of passed in addresses to ModuleEntryPoint
        dtm = self.currentProgram.getDataTypeManager()
        epType = dtm.getDataType("/UefiApplicationEntryPoint.h/functions/_ModuleEntryPoint")
        for index, entrypoint in enumerate(eps):
            self.updateFunctionDefinition(entrypoint, epType, "ModuleEntryPoint-{:d}".format(index))

    def findGlobalEfiPointers(self, func):
        # Identify instances of gST, gBS, EFI System Table, Image Handle
        hf = self.get_high_function(func)
        lPcodeOps = hf.getPcodeOps()
        while lPcodeOps.hasNext():
            var = lPcodeOps.next().getOutput()
            if var is not None and var.isAddress():
                varname = var.getHigh().getDataType().getName()
                if varname == "EFI_SYSTEM_TABLE *":
                    name = self.getNextLabel('gST')
                elif varname == "EFI_BOOT_SERVICES *":
                    name = self.getNextLabel('gBS')
                elif varname == "EFI_RUNTIME_SERVICES *":
                    name = self.getNextLabel('gRT')
                elif varname == "EFI_HANDLE":
                    name = self.getNextLabel('gImageHandle')
                else:
                    name = None
                if name is not None:
                    oAddr = self.toAddr(var.getOffset())
                    oDT = var.getHigh().getDataType()
                    if not artifacts().has_label(oAddr):
                        self.defineData(oAddr, oDT, name, None)
                        artifacts().add_label(oAddr, name, oDT)

    def identifyGBSProtocol(self, inFuncs=[]):
        # Identify and label protocols using gBS
        symbolName = "gBS"
        refs = self.getSymbolRefs(symbolName)
        refFuncs = self.getUniqueFuncts(refs)
        gbsFuncs = list(set(inFuncs + refFuncs))
        for func in gbsFuncs:
            hf = self.get_high_function(func)
            if hf is None:
                continue
            pcodeitr = hf.getPcodeOps()
            counter = 0
            while pcodeitr.hasNext():
                pcode = pcodeitr.next()
                if pcode.getOpcode() == PcodeOp.CALLIND:
                    if pcode.getInput(0).getHigh().getDataType().getName() == "EFI_LOCATE_PROTOCOL":
                        local = self.locateProtocol(pcode, 'gBS')
                        if local:
                            hfc = self.get_high_function(func)
                            pcodeitrCopy = hfc.getPcodeOps()
                            for i in range(counter):
                                if pcodeitrCopy.hasNext():
                                    pcodeitrCopy.next()
                                else:
                                    continue
                                self.propogateLocal(pcodeitrCopy, varnodeConverter(pcode.getInput(3)))
                    elif pcode.getInput(0).getHigh().getDataType().getName() == "EFI_INSTALL_PROTOCOL_INTERFACE":
                        self.installProtocol(pcode, 'gBS')
                counter += 1

    def identifySMST(self):
        # Identify and label SMST Instances
        # SMST must be identified after gBS
        symbolName = "gEFI_SMM_BASE2_PROTOCOL"
        symbolName = "gEFI_SMM_BASE2_PROTOCOL_GUID"
        symbolName2 = "EFI_SMM_BASE2_PROTOCOL_GUID"
        symbolName3 = "EFI_SMM_BASE2_PROTOCOL"
        refs = self.getSymbolRefs(symbolName)
        refs2 = self.getSymbolRefs(symbolName2, True)
        refs3 = self.getSymbolRefs(symbolName3)
        smstFuncs = self.getUniqueFuncts(refs + refs2 + refs3)
        for func in smstFuncs:
            hf = self.get_high_function(func)
            if hf is None:
                continue
            pcodeitr = hf.getPcodeOps()
            while pcodeitr.hasNext():
                pcode = pcodeitr.next()
                if pcode.getOpcode() == PcodeOp.CALLIND:
                    if pcode.getInput(0).getHigh().getDataType().getName() == "EFI_SMM_GET_SMST_LOCATION2":
                        if pcode.getNumInputs() - 1 == 2:
                            proto = pcode.getInput(2)
                            protoName = 'Smst'
                            #protoName = self.getNextLabel('Smst')
                            _, protoType = self.nameToProto("EFI_SMM_SYSTEM_TABLE2")
                            protoVar = varnodeConverter(proto)
                            self.labelvarNode(protoVar, protoName, protoType)

    def identifySMSTProtocol(self):
        # Identify and label protocols using SMST
        symbolName = "gSmst"
        refs = self.getSymbolRefs(symbolName)
        gbsFuncs = self.getUniqueFuncts(refs)
        for func in gbsFuncs:
            hf = self.get_high_function(func)
            if hf is None:
                continue
            pcodeitr = hf.getPcodeOps()
            while pcodeitr.hasNext():
                pcode = pcodeitr.next()
                if pcode.getOpcode() == PcodeOp.CALLIND:
                    if pcode.getInput(0).getHigh().getDataType().getName() == "EFI_LOCATE_PROTOCOL":
                        if pcode.getNumInputs() - 1 == 3:
                            self.locateProtocol(pcode, 'SMST')

    def identifySMSTHandlers(self):
        # Label Management Mode handlers
        toRegister = ["EFI_SMM_POWER_BUTTON_REGISTER2",
                      "EFI_SMM_SX_REGISTER2",
                      "EFI_SMM_SW_REGISTER2",
                      "EFI_SMM_PERIODIC_TIMER_REGISTER2",
                      "EFI_SMM_USB_REGISTER2",
                      "EFI_SMM_IO_TRAP_DISPATCH2_REGISTER",
                      "EFI_SMM_GPI_REGISTER2",
                      "EFI_MM__SW_REGISTER"]
        symbolName = "gSmst"
        refs = self.getSymbolRefs(symbolName)
        gbsFuncs = self.getUniqueFuncts(refs)
        for func in gbsFuncs:
            hf = self.get_high_function(func)
            if hf is None:
                continue
            pcodeitr = hf.getPcodeOps()
            counter = 0
            while pcodeitr.hasNext():
                pcode = pcodeitr.next()
                counter += 1
                if pcode.getOpcode() == PcodeOp.CALLIND:
                    if pcode.getInput(0).getHigh().getDataType().getName() in toRegister:
                        self.register2Handler(pcode, counter)
                    elif pcode.getInput(0).getHigh().getDataType().getName() == "EFI_SMM_INTERRUPT_REGISTER":
                        self.register2Handler(pcode)
                        self.labelVarnodeGuid(varnodeConverter(pcode.getInput(pcode.getNumInputs()-2)))
                    elif pcode.getInput(0).getHigh().getDataType().getName() == "EFI_SMM_REGISTER_PROTOCOL_NOTIFY":
                        self.labelVarnodeGuid(varnodeConverter(pcode.getInput(pcode.getNumInputs()-3)))
                        self.registerNotify(pcode)
                    elif pcode.getInput(0).getHigh().getDataType().getName() == "EFI_INSTALL_PROTOCOL_INTERFACE":
                        self.installProtocol(pcode, 'SMST')

    def installProtocol(self, pcode, origin):
        # label the installed protocol
        # pcode - pcode instruction containing install protocol
        # origin - where the protocol is installed (SMST or gBS)
        dtm = self.currentProgram.getDataTypeManager()
        handle = pcode.getInput(1)
        handleVar = varnodeConverter(handle)
        handleName = 'handle'
        handleType = dtm.getDataType("/UefiBaseType.h/EFI_HANDLE *")
        self.labelvarNode(handleVar, handleName, handleType)
        # label the guid
        guid = pcode.getInput(2)
        guidVar = varnodeConverter(guid)
        guidName = self.labelVarnodeGuid(guidVar)
        # label the interface
        iface = pcode.getInput(4)
        ifaceVar = varnodeConverter(iface)
        if guidName is not None and ifaceVar.offset != 0:
            ifaceName, ifaceType = self.nameToProto(guidName, True)
            if ifaceName is None:
                ifaceName = "unknown_{}".format(guidName[:-5])
            self.labelvarNode(ifaceVar, ifaceName, ifaceType, origin)

    def locateProtocol(self, pcode, origin):
        # Label the protocol used in locateProtocol
        # pcode - pcode instruction containing install protocol
        # origin - where the protocol is installed (SMST or gBS)
        guid = pcode.getInput(1)
        guidVar = varnodeConverter(guid)
        guidName = self.labelVarnodeGuid(guidVar)
        if guidName is not None:
            proto = pcode.getInput(3)
            protoVar = varnodeConverter(proto)
            protoName = None
            protoType = None
            protoName, protoType = self.nameToProto(guidName, True)
            if protoName is None:
                protoName = "unknown_{}".format(guidName[:-5])
            self.labelvarNode(protoVar, protoName, protoType, origin)
            # if protoVar is local need to propogate locally
            if protoVar.isLocal():
                return True
        return False

    def propogateLocal(self, pcodeitr, lvarnode):
        # propogae a local variable name within the rest of a function
        # pcodeitr - pcode iterator for the rest of the pcode within a function
        # lvarnode - local varnode
        # @TODO, check if the location already has a name, if so don't rename
        while pcodeitr.hasNext():
            pcode = pcodeitr.next()
            for num in range(pcode.getNumInputs()):
                try:
                    if lvarnode.offset == varnodeConverter(pcode.getInput(num)).offset:
                        if pcode.getOpcode() == PcodeOp.CALLIND:
                            return
                        if pcode.getOpcode() == PcodeOp.LOAD:
                            tmp = varnodeConverter(pcode.getOutput())
                            if tmp.isGlobal() and not artifacts().has_label(toAddr(tmp.offset)):
                                tmp.labelvarNode("fromGBS", None)
                        if pcode.getOpcode() == PcodeOp.COPY:
                            tmp = varnodeConverter(pcode.getOutput())
                            if tmp.isGlobal() and not artifacts().has_label(toAddr(tmp.offset)):
                                tmp.labelvarNode("fromGBS", None)
                        if pcode.getOpcode() == PcodeOp.INDIRECT:
                            tmp = varnodeConverter(pcode.getOutput())
                            if tmp.isGlobal() and not artifacts().has_label(toAddr(tmp.offset)):
                                tmp.labelvarNode("fromGBS", None)
                except:
                    print('Unable to propgate local')

    def getContextValue(self, vnode, inpcode, counter):
        # Get the context value of the Handler registered with EFI_SMM_SW_REGISTER2
        # vnode - pcode instance
        # inpcode - pcode where instance is contained
        # counter - number of iterations to attempt looking for the context value
        ret = None
        func = vnode.getHigh().getHighFunction().getFunction()
        cvnode = varnodeConverter(vnode)
        hf = self.get_high_function(func)
        pcodeitr = hf.getPcodeOps()
        loop_cnt = counter if counter else 1000
        while pcodeitr.hasNext():
            pcode = pcodeitr.next()
            if pcode.getOpcode() == PcodeOp.COPY:
                if pcode.getOutput().getOffset() == cvnode.offset:
                    ret = pcode.getInput(0).getAddress().getOffset()
            if counter is None and pcode.getOpcode() == inpcode.getOpcode() and pcode.getInput(1).getAddress().getOffset() == inpcode.getInput(1).getAddress().getOffset():
                break
            if loop_cnt == 0:
                break
            loop_cnt -= 1
        return ret

    def register2Handler(self, pcode, counter=None):
        # Label SMI Handler
        # pcode - pcode instance of where handler is located
        # counter - number of iterations to attempt looking for context value
        dtm = self.currentProgram.getDataTypeManager()
        tname = pcode.getInput(0).getHigh().getDataType().getName()
        if pcode.getInput(0).getHigh().getDataType().getName() == "EFI_SMM_SW_REGISTER2":
            vnode = pcode.getInput(pcode.getNumInputs()-2)
            contextNumber = self.getContextValue(vnode, pcode, counter)
            if contextNumber is not None:
                fname = "swSmiHandler_{:x}".format(contextNumber)
            else:
                fname = "swSmi_unknown"
        else:
            fname = self.getNextLabel("{}Handler".format(tname[8:-10].replace("_", "").replace("DISPATCH", "").lower()))
        vnode = varnodeConverter(pcode.getInput(pcode.getNumInputs()-3))
        faddr = self.toAddr(vnode.offset)
        fdef = dtm.getDataType("/PiMmCis.h/functions/EFI_MM_HANDLER_ENTRY_POINT")
        if vnode.isGlobal():
            self.updateFunctionDefinition(faddr, fdef, fname)

    def registerNotify(self, pcode):
        # Label SMI Notify Handler
        dtm = self.currentProgram.getDataTypeManager()
        fdef = dtm.getDataType("/PiMmCis.h/functions/EFI_MM_NOTIFY_FN")
        fname = self.getNextLabel("notifyHandler")
        vnode = varnodeConverter(pcode.getInput(pcode.getNumInputs()-2))
        faddr = self.toAddr(vnode.offset)
        if vnode.isGlobal():
            self.updateFunctionDefinition(faddr, fdef, fname)
            artifacts().add_notify(fname, artifacts().get_guid(self.toAddr(varnodeConverter(pcode.getInput(1)).offset))[0])

    def nameToProto(self, inName, isGuid=False, isPtr=True, TypeDefok=False):
        # Attempt to find a name within the datatypemanager
        # inName - Name to find
        # isGuid - is the Name a guid
        # isPtr - is the Name a pointer
        # TypeDefok - return typedef of name
        foundlist = []
        dtm = self.currentProgram.getDataTypeManager()
        if isGuid:
            tryName = inName[:-5]
        else:
            tryName = inName
        # attempt to find 4 variations _SMM_ -> _MM_ with * and without *
        if isPtr:
            dtm.findDataTypes("{} *".format(tryName.replace("_SMM_", "_MM_")), foundlist)
            dtm.findDataTypes("{} *".format(tryName), foundlist)
        else:
            dtm.findDataTypes("{}".format(tryName.replace("_SMM_", "_MM_")), foundlist)
            dtm.findDataTypes("{}".format(tryName), foundlist)
        if foundlist:
            if not TypeDefok and str(foundlist[0]).find('typedef') >= 0:
                if len(foundlist) == 1:
                    searchstr = "{}".format(str(foundlist[0]).split(' ')[2])
                    return self.nameToProto(searchstr, isPtr=isPtr)
                else:
                    return (tryName, foundlist[1])
            else:
                return (tryName, foundlist[0])
        else:
            return (None, None)

    def labelVarnodeGuid(self, guidVar):
        # label a GUID
        # guidVar - Varnode guid
        dtm = self.currentProgram.getDataTypeManager()
        varAddr = self.toAddr(guidVar.offset)
        guidName, _ = artifacts().get_guid(varAddr)
        if not artifacts().has_label(varAddr):
            if guidName is None:
                if guidVar.isGlobal():
                    try:
                        guidbytes = self.getBytes(varAddr, 16)
                        memGuid = UUID(bytes_le=struct.unpack('16s', guidbytes)[0])
                    except MemoryAccessException:
                        memGuid = UUID("FFFFFFFF-FFFF-FFFF-FFFF-FFFFFFFFFFFF")
                    guidName = "unknown_{}".format(str(memGuid)[:13])
                elif guidVar.isLocal():
                    guidName = self.getLocalGuidValue(guidVar)
            efiGuidType = dtm.getDataType("/UefiBaseType.h/EFI_GUID")
            self.labelvarNode(guidVar, guidName, efiGuidType)
        return guidName

    def getLocalGuidValue(self, guidVar):
        # Get the GUID value from a local variable
        # guidVar - VarNode that contains GUID
        guidList = []
        func = guidVar.varnode.getHigh().getHighFunction().getFunction()
        hf = self.get_high_function(func)
        pcodeitr = hf.getPcodeOps()
        while pcodeitr.hasNext():
            pcode = pcodeitr.next()
            if pcode.getOpcode() == PcodeOp.COPY:
                if pcode.getOutput().getOffset() in [guidVar.offset, guidVar.addr.add(4).getOffset(), guidVar.addr.add(8).getOffset(), guidVar.addr.add(12).getOffset()]:
                    guidList.append(pcode.getInput(0).getAddress().getOffset())
        guidstr = ""
        for i in guidList:
            guidstr += "{0:08X}".format(i)
        guidstr = convertGuidStr(guidstr)
        if len(guidstr) != 32:
            guidName = None
        else:
            try:
                guid = UUID(guidstr)
                guidName = guids().getGuidName(guid)
            except ValueError:
                guidName = None
        if guidName is None:
            guidName = "unknownGuid"
        return guidName

    def labelvarNode(self, node, name, dataType, origin=None):
        # label a varnode
        # node - node to label
        # datatype - what datatype to assign the node
        # origin - where the label originated
        if node.isGlobal():
            self._labelGlobalvarNode(node, name, dataType, origin)
        elif node.isLocal():
            self._labelLocalvarNode(node, name, dataType)
        else:
            print("unable to label unknown varnode address location")

    def _labelGlobalvarNode(self, node, name, dataType, origin):
        # label a  global varnode
        # node - node to label
        # datatype - what datatype to assign the node
        # origin - where the label originated
        varName = self.getNextLabel('g{}'.format(name))
        if not artifacts().has_protocol_global(self.toAddr(node.offset)):
            self.defineData(self.toAddr(node.offset), dataType, varName)
            artifacts().add_protocol_global(self.toAddr(node.offset), varName, dataType, origin)
        else:
            print("Address {} already has a label".format(node.addr))

    def _labelLocalvarNode(self, node, name, dataType):
        # label a local varnode
        # node - node to label
        # datatype - what datatype to assign the node
        varName = self.getNextLabel(name)
        func = node.varnode.getHigh().getHighFunction().getFunction()
        funcVars = func.getAllVariables()
        labelVar = None
        for variable in funcVars:
            if node.spaceId == AddressSpace.TYPE_REGISTER:
                if variable.isRegisterVariable() and variable.getRegister().getOffset() == node.addr:
                    labelVar = variable
                    break
            elif node.spaceId == AddressSpace.TYPE_VARIABLE:
                if variable.isStackVariable() and variable.getStackOffset() == node.offset:
                    labelVar = variable
                    break
        if labelVar is not None:
            node.defineVar(labelVar, dataType, varName)
            artifacts().add_protocol_local(func, labelVar, varName, dataType)

    def labelGuids(self):
        # search memory from minaddress to maxaddress looking for
        # potential guid matches
        dtm = self.currentProgram.getDataTypeManager()
        efiGuidType = dtm.getDataType("/UefiBaseType.h/EFI_GUID")

        start = self.currentProgram.getMinAddress()
        binSize = self.currentProgram.getMaxAddress().getOffset() - start.getOffset()

        for addrPlus in range(0, binSize, 4):
            addr = start.add(addrPlus)
            try:
                rawGuid = self.getBytes(addr, 16)
            except MemoryAccessException:
                continue
            memGuid = UUID(bytes_le=struct.unpack('16s', rawGuid)[0])
            if memGuid == UUID("00000000-0000-0000-0000-000000000000"):
                continue
            guidName = guids().getGuidName(memGuid)
            if guidName is not None:
                self.defineData(addr, efiGuidType, guidName, None)
                artifacts().add_guid(addr, guidName, memGuid)

    def functionVariableUse(self, func):
        """Return dictionary of GET_VARIABLE and SET_VARIABLE results within a function"""
        """SET_VARIABLE will return Name, GUID, size, and if the Data is concrete"""
        """GET_VARIBLE will return Name, GUID, and if the return value is checked"""
        # func - function to check
        hf = self.get_high_function(func)
        if hf is None:
            return {}
        pcodeitr = hf.getPcodeOps()
        res = {}
        counter = 0
        while pcodeitr.hasNext():
            pcode = pcodeitr.next()
            if pcode.getOpcode() == PcodeOp.CALLIND:
                tname = pcode.getInput(0).getHigh().getDataType().getName()
                if tname.find('ET_VARIABLE') > 0:
                    if tname not in res.keys():
                        res[tname] = []
                    nnode = varnodeConverter(pcode.getInput(1))
                    gnode = varnodeConverter(pcode.getInput(2))
                    gname = None
                    gguid = None
                    if nnode.isGlobal():
                        gname = self.getDataAt(self.toAddr(nnode.offset))
                    else:
                        gname = self.getContextValue(pcode.getInput(1), pcode, counter)
                        if gname is not None and not self.getDataAt(self.toAddr(gname)):
                            clearsize = self.findUnicodeSize(self.toAddr(gname))
                            if clearsize is not None:
                                listing = self.currentProgram.getListing()
                                listing.clearCodeUnits(self.toAddr(gname), self.toAddr(gname).add(clearsize), True)
                                self.createData(self.toAddr(gname), UnicodeDataType())
                        if gname is not None:
                            gname = self.getDataAt(self.toAddr(gname))
                    if gnode.isGlobal():
                        gguid = self.getSymbolAt(self.toAddr(gnode.offset))
                    else:
                        gguid = self.getLocalGuidValue(gnode)
                    if tname.find('SET_VARIABLE') > 0:
                        snode = varnodeConverter(pcode.getInput(4))
                        concreateSize = (pcode.getInput(4).getSpace() & AddressSpace.ID_TYPE_MASK) == AddressSpace.TYPE_CONSTANT
                        if snode.isGlobal():
                            ssize = snode.offset
                        else:
                            ssize = self.getContextValue(pcode.getInput(4), pcode, counter)
                        res[tname].append((gname, gguid, ssize, concreateSize))
                    elif tname.find('GET_VARIABLE') > 0:
                        rtchecked = False
                        if pcode.getOutput() is not None:
                            rnode = varnodeConverter(pcode.getOutput())
                            itercpy = hf.getPcodeOps()
                            for _ in range(counter+1):
                                itercpy.next()
                            rtchecked = self.returnChecked(itercpy, rnode)
                        res[tname].append((gname, gguid, rtchecked))
            counter += 1
        return res

    def findUnicodeSize(self, Addr):
        try:
            mbytes = self.getBytes(Addr, 64)
            for i in range(0, 63, 2):
                if mbytes[i] == 0 and mbytes[i + 1] == 0:
                    break
            return i + 1
        except MemoryAccessException:
            return None

    def returnChecked(self, pcodeitr, varnode):
        res = False
        while pcodeitr.hasNext():
            pcode = pcodeitr.next()
            if pcode.getOutput() and varnodeConverter(pcode.getOutput()).offset == varnode.offset:
                break
            for num in range(pcode.getNumInputs()):
                if varnode.offset == varnodeConverter(pcode.getInput(num)).offset:
                    if pcode.getOpcode() < 17 and pcode.getOpcode() > 10:
                        res = True
                        break
        return res

UEFISurveyor/EFI_functs.py
dcg-red-team/UEFI_Surveyor
# Copyright 2023 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# @category UEFISurveyor.internal

import json
from ghidra.app.decompiler.component import DecompilerUtils
from EFI_functs import EFIUtils
from logger import logger
from artifacts import artifacts
from ghidra.feature.fid.service import FidService


class CallOuts:
    def __init__(self, name, func_list):
        self.name = name
        self.funcs = func_list
        self.callouts = []

    def get_funcs(self):
        return self.funcs

    def add_callout(self, func, handler, path):
        self.callouts.append((func, handler, path))

    def log_results(self):
        if not self.callouts:
            logger().log("Did not find any callouts in from {}".format(self.name))
            logger().log('')
        else:
            for (func, handler, path) in self.callouts:
                logger().log("Potential Callout:")
                logger().log('{} has callpath {} to {} with {}'.format(handler, path, func, self.name))


class EFIAnalytics(EFIUtils):

    def __init__(self, currentProgram):
        super(EFIAnalytics, self).__init__(currentProgram)

    def identifyCallouts(self):
        logger().log('Searching for potential callouts')
        # Gather the functions that contain gBS and gRS
        refs = []
        gbsrefs = self.getSymbolRefs("gBS")
        refs.append(CallOuts('gBS', self.getUniqueFuncts(gbsrefs)))

        grsrefs = self.getSymbolRefs("gRS")
        refs.append(CallOuts('gRS', self.getUniqueFuncts(grsrefs)))

        # Gather any artifacts that came from gRT or gBS
        for key in artifacts().protocols.keys():
            name, _, origin = artifacts().protocols[key]
            if origin == 'gBS' and name not in ['gBS', 'gRS']:
                globalrefs = self.getSymbolRefs(name, True)
                refs.append(CallOuts(name, self.getUniqueFuncts(globalrefs)))

        # Enumerate the SMIHandlers
        # search for a callpath that goes to a function with the handler
        funcMan = self.currentProgram.getFunctionManager()
        for func in [func for func in funcMan.getFunctions(True) if func.getName().find("Handler") >= 0]:
            for ref in refs:
                for gfunc in ref.get_funcs():
                    emptylist = []
                    fnPath = self.isFunctionPath(func, [gfunc], emptylist)
                    if fnPath:
                        ref.add_callout(gfunc, func, fnPath)
        for ref in refs:
            ref.log_results()
        logger().log('')

    # Crude check Within handler functions check whether commbuffer is dereferenced multiple times
    # The theory is that the data should be copied and worked on by a local buffer if there are multiple
    # times that it is dereferenced that may indicsate that a potential TOCTOU window
    # @TODO add savestate registers
    def identifyTOCTOU(self):
        funcMan = self.currentProgram.getFunctionManager()
        for func in [func for func in funcMan.getFunctions(True) if func.getName().find("swSmiHandler") >= 0]:
            hf = self.get_high_function(func)
            lsm = hf.getLocalSymbolMap()
            symbols = lsm.getSymbols()

            for symbol in symbols:
                if symbol.getName() != "CommBuffer":
                    continue
                hv = symbol.getHighVariable()
                varnode = hv.getRepresentative()
                fswo = DecompilerUtils.getForwardSliceToPCodeOps(varnode)
                count = 0
                for f in fswo:
                    if "LOAD" in str(f):
                        count += 1
                logger().log("{} CommBuffer is dereferenced {} times".format(func, count))

    # Find the uses of GetVariable and SetVariable within a function
    def identifyVariableUses(self):
        logger().log('Searching for (G/S)ET_VARIABLE:')
        funcMan = self.currentProgram.getFunctionManager()
        found = False
        for func in funcMan.getFunctions(True):
            results = self.functionVariableUse(func)
            if results:
                found = True
                logger().log('{}:'.format(func))
                for key in results.keys():
                    logger().log("  {}:".format(key))
                    for element in results[key]:
                        logger().log('    {}'.format(element))
        if not found:
            logger().log('Did not find any instances of (G/S)ET_VARIABLE')
        logger().log('')

    def getFunctionHashes(self, programFile):
        funcsDict = {}
        service = FidService()
        logger().log('Function Hashes')
        funcMan = self.currentProgram.getFunctionManager()
        for func in funcMan.getFunctions(True):
            hf = service.hashFunction(func)
            if hf is None:
                continue
            logger().log('{} - {}'.format(func, hf))
            name = func.getName()
            if name in funcsDict:
                name = self.getUniqueName(funcsDict.keys(), name)
            funcsDict[name] = str(hf)
        funcsJson = json.dumps(funcsDict, indent=4)
        outfile = '{}.json'.format(programFile)
        with open(outfile, 'w') as jsonout:
            jsonout.write(funcsJson)

    def getUniqueName(self, keys, name):
        newname = name
        for i in range(100):
            newname = '{}_{:d}'.format(name, i)
            if newname not in keys:
                break
        return newname

UEFISurveyor/analytics.py
dcg-red-team/UEFI_Surveyor
# Copyright 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# This script identifies module entry point and labels
# efi specific functions, protocols, and guids
# @category UEFISurveyor.headless

from ghidra.util.task import ConsoleTaskMonitor
import ghidra.util.Msg as Msg
from ghidra.program.disassemble import Disassembler
from ghidra.app.plugin.core.analysis import AutoAnalysisManager


from logger import logger
from artifacts import artifacts
from guids import guids
from EFI_functs import EFIUtils
from analytics import EFIAnalytics

import sys


if __name__ == "__main__":
    file = getProgramFile()
    logger().enableLogFile('{}.log'.format(file))
    logger().log("Attempting to Analyze {}".format(file))

    # initialize GhidraUtils class
    EFIUtil = EFIUtils(currentProgram)
    EFIAnalytic = EFIAnalytics(currentProgram)

    # load guid file
    args = getScriptArgs()
    logger().log("Args passed in: {}".format(args))
    if len(args) > 0:
        guidFile = args[0]
    else:
        logger().log('Missing Guid DB file!')
        sys.exit()
    guids().loadGuidFile(guidFile)

    # label Guids with the file
    EFIUtil.labelGuids()

    # Search for Entry Points
    eps = EFIUtil.findModuleEntries()

    # ensure Decompiler Parameter ID is enabled
    setAnalysisOption(currentProgram, "Decompiler Parameter ID", "true")

    # Disassemble the entry points
    ctMon = ConsoleTaskMonitor()
    dis = Disassembler.getDisassembler(currentProgram, ctMon, None)
    for entry in eps:
        disAddrs = dis.disassemble(entry, None)
        AutoAnalysisManager.getAnalysisManager(currentProgram).codeDefined(disAddrs)
        AutoAnalysisManager.getAnalysisManager(currentProgram).waitForAnalysis(None, ctMon)

    # label entry points
    EFIUtil.labelModuleEntryPoints(eps)

    # propogate variables from entrypoint to find global tables
    propFuncs = []
    for entry in eps:
        entryFunc = getFunctionAt(entry)
        propFuncs += EFIUtil.propogateFunctionVariables(entryFunc, callonExit=EFIUtil.findGlobalEfiPointers)

    # label instances of locate protocol
    EFIUtil.identifyGBSProtocol(propFuncs)

    # label instances of SMST
    EFIUtil.identifySMST()

    # label insances where SMST uses locate protocol
    EFIUtil.identifySMSTProtocol()

    # label Management Mode handlers
    EFIUtil.identifySMSTHandlers()

    # log what is found
    artifacts().logArtifacts()

    # attempt to find callouts
    EFIAnalytic.identifyCallouts()

    # attempt to find TOCTOU instances
    EFIAnalytic.identifyTOCTOU()

    # attempt to find variable functionality
    EFIAnalytic.identifyVariableUses()

    # log the functions and hashes
    EFIAnalytic.getFunctionHashes(file)

    logger().closeLogFile()

UEFISurveyor/analyze_post.py
dcg-red-team/UEFI_Surveyor
# Copyright 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# @category UEFISurveyor.headless

from ghidra_funcs import loadGDTFile
import ghidra.util.Msg as Msg
from logger import logger
import sys

prefix_text = "[Pre]"


def output(text):
    logger().log('{} {}'.format(prefix_text, text))


if __name__ == "__main__":
    output("Welcome to the Analyse Pre Script")
    try:
        currentProgram.setImageBase(toAddr(0x80000000), True)
    except Exception as e:
        Msg.error("EFI Initialization", "Problems moving base address\n{}".format(e))

    # Set the write attribute on the .text section
    block = getMemoryBlock('.text')
    if block is not None:
        block.setWrite(True)

    # load gdt file
    args = getScriptArgs()
    output("Args passed in: {}".format(args))
    if len(args) > 0:
        gdtPath = args[0]
    else:
        logger().log('Missing Guid DB file!')
        sys.exit()

    output("Path is {}".format(gdtPath))

    loadGDTFile(gdtPath, currentProgram)

UEFISurveyor/analyze_pre.py
dcg-red-team/UEFI_Surveyor
# Copyright 2023 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# @category UEFISurveyor.internal

from logger import logger


class Artifacts:
    # Class that keeps track of changes made to EFI file
    # Includes guids, protocols, labels, and functions
    def __init__(self):
        self.guids = {}
        self.protocols = {}
        self.localProtocols = {}
        self.labels = {}
        self.functions = {}
        self.notify = {}

    def add_guid(self, address, name, bytes):
        if address in self.guids.keys():
            logger().log("Guid already defined at address {} - {}".format(address, self.guids[address]))
        self.guids[address] = (name, bytes)

    def get_guid(self, address):
        ret = (None, None)
        if address in self.guids.keys():
            ret = self.guids[address]
        return ret

    def add_protocol_global(self, address, name, datatype, origin):
        if address in self.protocols.keys():
            logger().log("protocol already defined for address {} - {}".format(address, self.protocols[address]))
        self.protocols[address] = (name, datatype, origin)

    def has_protocol_global(self, address):
        if address in self.protocols.keys():
            return True
        return False

    def add_protocol_local(self, function, location, name, datatype):
        if name in self.localProtocols.keys():
            logger().log("duplicate key found adding local protocol")
        self.localProtocols[name] = (function, location, datatype)

    def add_label(self, address, name, datatype):
        if address in self.labels.keys():
            logger().log("duplicate key found adding label")
        self.labels[address] = (name, datatype)

    def has_label(self, address):
        if address in self.labels.keys():
            return True
        else:
            return False

    def add_function(self, address, name):
        if address in self.functions:
            logger().log('Renaming funciton at {}'.format(address))
        self.functions[address] = name

    def add_notify(self, name, event):
        self.notify[name] = event

    def logArtifacts(self):
        logger().log('')
        logger().log('Identified the following Guids')
        for key in self.guids.keys():
            logger().log("{} - {}".format(self.guids[key][0], key))
        logger().log('')
        logger().log('Identified the following protocols')
        for key in self.protocols.keys():
            logger().log("{} - {}".format(key, self.protocols[key]))
        logger().log('')
        for key in self.localProtocols.keys():
            logger().log("{} - {}".format(key, self.localProtocols[key]))
        logger().log('')
        logger().log('Applied the following labels:')
        for key in self.labels.keys():
            logger().log("{} - {}".format(key, self.labels[key]))
        logger().log('')
        logger().log('Identified the following functions:')
        for key in self.functions.keys():
            logger().log("{} - {}".format(key, self.functions[key]))
        logger().log('')
        logger().log('Notify Handlers registered events:')
        for key in self.notify.keys():
            logger().log('{} - {}'.format(key, self.notify[key]))
        logger().log('')


_Artifacts = Artifacts()


def artifacts():
    return _Artifacts

UEFISurveyor/artifacts.py
dcg-red-team/UEFI_Surveyor
# Copyright 2023 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# @category UEFISurveyor.internal

from ghidra.app.decompiler import DecompileOptions, DecompInterface
from ghidra.util.task import ConsoleTaskMonitor
import ghidra.program.model.symbol.SourceType as SourceType
from ghidra.program.model.listing import ParameterImpl
from ghidra.program.model.listing import ReturnParameterImpl
import ghidra.util.Msg as Msg
from ghidra.program.model.data import GenericCallingConvention
from ghidra.app.decompiler.component import DecompilerUtils
from ghidra.program.model.pcode import PcodeOp
from ghidra.program.model.address import AddressSpace
import ghidra.program.flatapi.FlatProgramAPI as FlatAPI
import java.io.File as File
from ghidra.program.model.listing.Function import FunctionUpdateType
import ghidra.program.model.data.DataTypeConflictHandler as ConflictHandler
from ghidra.program.model.util import CodeUnitInsertionException

from artifacts import artifacts

import os
import struct


class GhidraUtils(FlatAPI):
    # Generic functions to expand GhidraAPIs for use within UEFISurveyor
    def __init__(self, currentProgram):
        super(GhidraUtils, self).__init__(currentProgram)

    def defineData(self, address, datatype=None, name=None, comment=None):
        # Defines data at an address
        # address - where to label the data
        # datatype - datatype to apply to address
        # name - label to apply to address
        # comment - comment to apply at address
        if address == self.toAddr(0):
            print("unable to DefineData at address 0")
            return False
        if datatype is not None:
            Msg.info("Defining Data", "Overwriting address {} with datatype {}".format(address, datatype))
            listing = self.currentProgram.getListing()
            clrLen = datatype.getLength() - 1
            if clrLen > 0:
                listing.clearCodeUnits(address, address.add(clrLen), True)
            try:
                self.createData(address, datatype)
            except CodeUnitInsertionException:
                pass
        if name is not None:
            self.createLabel(address, name, True, SourceType.USER_DEFINED)
        if comment is not None:
            self.setPlateComment(address, comment)
        return True

    def findModuleEntries(self):
        # Find module entry point of UEFI file
        # Note - if using individual files can rely on PE parsing and labeling
        entrypoints = []
        start = self.currentProgram.getMinAddress()
        maxAddr = self.currentProgram.getMaxAddress()
        # search for 'MZ' within loaded files
        mz_header = self.findBytes(start, 'MZ', 1)
        if mz_header and mz_header[0] == start:
            # search for 'PE' within 'MZ' entries
            for entry in mz_header:
                # check for 'PE'
                peHeader = self.findBytes(entry, 'PE', 1)
                if not peHeader:
                    continue
                peAddr = [i for i in peHeader][0]
                # @TODO verify this holds true in all cases
                if peAddr and peAddr.subtract(entry) < 0x100:
                    # Get EntryPoint address
                    ep = entry.add(struct.unpack('<I', self.getBytes(peAddr.add(0x28), 4))[0])
                    # Get Image Size
                    imageSize = entry.add(struct.unpack('<I', self.getBytes(peAddr.add(0x50), 4))[0])
                    # Sanity check to see if entrypoint is less than imageSize
                    if ep < imageSize and ep < maxAddr:
                        entrypoints.append(ep)
        # search for terse executable
        if not entrypoints:
            vz_header = self.findBytes(start, 'VZ', 100)
            for entry in vz_header:
                ep = entry.add(struct.unpack('<I', self.getBytes(entry.add(0x8), 4))[0])
                if ep < maxAddr:
                    entrypoints.append(ep)
        return entrypoints

    def propogateFunctionVariables(self, func, params=[], visited=[], callonExit=None):
        # propogae function variable to new functions
        # func - function to propogate variable from
        # params - specific parameters to propogate
        # visited - functions visited
        # callonExit - function to Call before exiting
        visited.append(func)
        hf = self.get_high_function(func)
        if hf is None:
            return visited
        hfProto = hf.getFunctionPrototype()
        ret = {}
        # for each param see if it's used in a call
        for i in range(hfProto.getNumParams()):
            hfParam = hfProto.getParam(i)
            paramName = hfParam.getName()
            if not params or paramName in params:
                paramLoc = hfParam.getStorage()
                paramDT = hfParam.getDataType()
                hv = hfParam.getHighVariable()
                varnode = hv.getRepresentative()
                fswo = DecompilerUtils.getForwardSliceToPCodeOps(varnode)
                for f in fswo:
                    if f.getOpcode() == PcodeOp.CALL:
                        toFunc = f.getInput(0).getAddress()
                        if toFunc in ret.keys():
                            ret[toFunc].append((paramName, paramLoc, paramDT))
                        else:
                            ret[toFunc] = [(paramName, paramLoc, paramDT)]
        # gather params
        for key in ret.keys():
            pFunc = self.getFunctionAt(key)
            if pFunc is None:
                continue
            params, funcDef = self.createParamList(pFunc, ret[key])
        # call change func definition
            self.updateFunctionDefinition(key, funcDef, pFunc.getName(), params)
        # call recursive for only params changed
            if pFunc not in visited:
                forwardparams = []
                for name, _, _ in ret[key]:
                    forwardparams.append(name)
                self.propogateFunctionVariables(pFunc, forwardparams, visited, callonExit)
        if callonExit is not None:
            callonExit(func)
        return visited

    def updateFunctionDefinition(self, funcAddr, funcDef, funcName=None, funcParams=None):
        # update a function definition
        # funcAddr - Address of function
        # funcDef - Original funciton definition
        # funcName - Name for function
        # funcParams - New function Parameters
        if funcAddr == self.toAddr(0):
            return False
        if funcParams is None:
            parameterDefinitions = funcDef.getArguments()
            parameters = []
            for d in parameterDefinitions:
                parameters.append(ParameterImpl(d.getName(), d.getDataType(), self.currentProgram))
        else:
            parameters = funcParams

        if funcName is None:
            funcName = funcDef.getName()
        returnType = ReturnParameterImpl(funcDef.getReturnType(), self.currentProgram)
        func = self.getFunctionAt(funcAddr)
        if func is None:
            func = self.createFunction(funcAddr, funcName)
            if func is None:
                Msg.error("updateFunction", "Unable to create function: {} at address{}".format(funcName, funcName))
        else:
            func.setName(funcName, SourceType.USER_DEFINED)
        if func is not None:
            if self.currentProgram.getMetadata().get('Address Size') == '64':
                func.updateFunction(GenericCallingConvention.fastcall.toString(), returnType, parameters, FunctionUpdateType.DYNAMIC_STORAGE_FORMAL_PARAMS, False, SourceType.USER_DEFINED)
            else:
                func.updateFunction(GenericCallingConvention.cdecl.toString(), returnType, parameters, FunctionUpdateType.DYNAMIC_STORAGE_FORMAL_PARAMS, False, SourceType.USER_DEFINED)
        artifacts().add_function(funcAddr, funcName)
        return True

    def get_markup_function(self, func):
        # get a function markup
        # func - function name
        options = DecompileOptions()
        monitor = ConsoleTaskMonitor()
        ifc = DecompInterface()
        ifc.setOptions(options)
        ifc.openProgram(self.currentProgram)
        res = ifc.decompileFunction(func, 60, monitor)
        tokenGroup = res.getCCodeMarkup()
        return tokenGroup

    def findDATinMarkup(self, tokenGroup):
        # Find DAT entries in Markup
        # tokenGroup - markup
        ret = set()
        tgString = tokenGroup.toString()
        loc = 0
        while loc != -1:
            loc = tgString.find('DAT_', loc)
            if loc != -1:
                ret.add(tgString[loc: loc+12])
                loc += 1
        return list(ret)

    def get_high_function(self, func):
        # Return Ghidra HighFunction class
        # func - function to decompile
        options = DecompileOptions()
        monitor = ConsoleTaskMonitor()
        ifc = DecompInterface()
        ifc.setOptions(options)
        ifc.openProgram(self.currentProgram)
        res = ifc.decompileFunction(func, 60, monitor)
        high = res.getHighFunction()
        return high

    def createParamList(self, func, changeParams):
        # Create a list of function parameters
        # func - function name
        # changeParams - Parameters to change from original
        outparams = []
        hf = self.get_high_function(func)
        hfProto = hf.getFunctionPrototype()
        for i in range(hfProto.getNumParams()):
            hfParam = hfProto.getParam(i)
            paramName = hfParam.getName()
            paramLoc = hfParam.getStorage()
            paramDT = hfParam.getDataType()
            for name, loc, DT in changeParams:
                if loc == paramLoc:
                    paramName = name
                    paramDT = DT
            outparams.append(ParameterImpl(paramName, paramDT, self.currentProgram))
        return (outparams, hfProto)

    def getSymbolRefs(self, symbolName, unique=False):
        """Return a list of references to a symbol"""
        # symbolName - symbol Name to get references of
        # unique - Attempt to find all references matching the name
        ret = []
        if unique:
            limit = 1
        else:
            limit = 100
        symbolTable = self.currentProgram.getSymbolTable()
        for index in range(limit):
            if not unique:
                newlabel = "{}_{:d}".format(symbolName, index)
            else:
                newlabel = symbolName
            found = [i for i in symbolTable.getSymbols(newlabel)]
            if found:
                for sym in found:
                    for ref in sym.getReferences():
                        refAddr = ref.getFromAddress()
                        if refAddr not in ret and refAddr is not None:
                            ret.append(refAddr)
            if not found:
                return ret
        return ret

    def getUniqueFuncts(self, addrlist):
        """Return a list of Unique functions from a list of Addresses"""
        # addrlist - list of addresses
        listing = self.currentProgram.getListing()
        ret = []
        for addr in addrlist:
            functionWithReference = listing.getFunctionContaining(addr)
            if functionWithReference not in ret and functionWithReference is not None:
                ret.append(functionWithReference)
        return ret

    def getNextLabel(self, label):
        """Generate Unique label from input"""
        # label - name of label
        symbolTable = self.currentProgram.getSymbolTable()
        index = 0
        while True:
            newlabel = "{}_{:d}".format(label, index)
            found = [i for i in symbolTable.getSymbols(newlabel)]
            if not found:
                return newlabel
            index += 1
            if index > 100:
                return label

    def isFunctionPath(self, start, goal, gSearchList=[], path=[]):
        """Return paths if a function can be reach a list of functions"""
        # start - initial function
        # goal - function to find a path into
        # gSearchList - functions searched
        # path - list of functions in path
        if not goal:
            return []
        if goal == start:
            path.append([start])
            return path
        if start not in gSearchList:
            gSearchList.append(start)
            for called in start.getCalledFunctions(self.getMonitor()):
                res = self.isFunctionPath(called, goal, gSearchList, [])
                if res:
                    for found in res:
                        found.append(start)
                        path.append(found)
        return path


class varnodeConverter(object):
    # class to hold varnode, address, AddressSpace type, and address offset
    # used to find and label individual elements within pcode
    def __init__(self, varnode):
        self.varnode = varnode
        self.addr, self.spaceId, self.offset = self.findvarNodeAddr()

    def findvarNodeAddr(self):
        # identify address, address space type and offset of a pcode element
        offset = self.varnode.getAddress().getOffset()
        addr = self.varnode.getAddress()
        retType = self.varnode.getSpace() & AddressSpace.ID_TYPE_MASK
        if self.varnode.isUnique():
            newvarNode = self.varnode.getDef()
            var0 = newvarNode.getInput(0)
            if var0.isUnique():
                var = varnodeConverter(var0)
                return var.findvarNodeAddr()
            if var0.isRegister():
                if var0.getAddress().getOffset() >= 0x20 and newvarNode.getNumInputs() > 1:
                    retType = AddressSpace.TYPE_VARIABLE
                    offset = newvarNode.getInput(1).getAddress().getOffset()
                    addr = newvarNode.getInput(1).getAddress()
                elif var0.getAddress().getOffset() >= 0x20:
                    retType = AddressSpace.TYPE_VARIABLE
                    addr = var0.getAddress()
                    offset = var0.getAddress().getOffset()
                else:
                    retType = var0.getSpace() & AddressSpace.ID_TYPE_MASK
                    addr = var0.getAddress()
                    offset = var0.getAddress().getOffset()
            elif var0.isConstant():
                var1 = newvarNode.getInput(1)
                if var1 is not None:
                    retType = var1.getSpace() & AddressSpace.ID_TYPE_MASK
                    offset = var1.getAddress().getOffset()
                    addr = var1.getAddress()
                else:
                    retType = var0.getSpace() & AddressSpace.ID_TYPE_MASK
                    offset = var0.getAddress().getOffset()
                    addr = var0.getAddress()
        return addr, retType, offset

    def isGlobal(self):
        """return whether varnode is global to function"""
        if self.spaceId in [AddressSpace.TYPE_CONSTANT, AddressSpace.TYPE_RAM]:
            return True
        return False

    def isLocal(self):
        """return whether varnode is local to function"""
        if self.spaceId in [AddressSpace.TYPE_REGISTER, AddressSpace.TYPE_VARIABLE]:
            return True
        return False

    def defineVar(self, var, dataType=None, name=None):
        """Set a new Name or dataType for a variable"""
        if name is not None:
            var.setName(name, SourceType.USER_DEFINED)
        if dataType is not None:
            var.setDataType(dataType, False, True, SourceType.USER_DEFINED)


def loadGDTFile(gdtPath, currentProgram):
    mflatAPI = FlatAPI(currentProgram)
    dtm = currentProgram.getDataTypeManager()
    gdtFile = File(gdtPath)
    newDTM = mflatAPI.openDataTypeArchive(gdtFile, True)
    dtl = []
    newDTM.getAllDataTypes(dtl)
    kh = ConflictHandler.KEEP_HANDLER
    dtm.addDataTypes(dtl, kh, ConsoleTaskMonitor())

UEFISurveyor/ghidra_funcs.py
dcg-red-team/UEFI_Surveyor
# Copyright 2023 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# @category UEFISurveyor

from ghidra.util.task import ConsoleTaskMonitor
from ghidra.app.decompiler.component import DecompilerUtils
import ghidra.util.Msg as Msg
from ghidra.program.disassemble import Disassembler
from ghidra.app.plugin.core.analysis import AutoAnalysisManager
from ghidra.program.model.mem import MemoryAccessException

from ghidra_funcs import loadGDTFile
from logger import logger
from artifacts import artifacts
from guids import guids
from EFI_functs import EFIUtils
from analytics import EFIAnalytics

if __name__ == "__main__":
    logger().log("Attempting to Analyze EFI file")
    # Set Filenames for GUIDDB and GDT file
    guidFile = '/home/bh/devel/Demo/Tools/guidFinder/tmp_guid_db'
    gdtPath = '/home/bh/devel/Demo/Tools/prfGenerator/gdt_files/uefi_x64.gdt'

    # reset image base
    try:
        currentProgram.setImageBase(toAddr(0x80000000), True)
    except Exception as e:
        Msg.error("EFI Initialization", "Problems moving base address\n{}".format(e))

    # Set the write attribute on the .text section
    block = getMemoryBlock('.text')
    if block:
        block.setWrite(True)

    # load the definitions file
    loadGDTFile(gdtPath, currentProgram)

    # load the guid file
    guids().loadGuidFile(guidFile)

    # initialize GhidraUtils class
    EFIUtil = EFIUtils(currentProgram)
    EFIAnalytic = EFIAnalytics(currentProgram)

    #label Guids
    EFIUtil.labelGuids()

    # Search for Entry Points
    eps = EFIUtil.findModuleEntries()

    # ensure Decompiler Parameter ID is enabled
    setAnalysisOption(currentProgram, "Decompiler Parameter ID", "true")

    # Disassemble the entry points
    ctMon = ConsoleTaskMonitor()
    dis = Disassembler.getDisassembler(currentProgram, ctMon, None)
    for entry in eps:
        disAddrs = dis.disassemble(entry, None)
        AutoAnalysisManager.getAnalysisManager(currentProgram).codeDefined(disAddrs)
        AutoAnalysisManager.getAnalysisManager(currentProgram).waitForAnalysis(None, ctMon)

    # label entry points
    EFIUtil.labelModuleEntryPoints(eps)

    # propogate variables from entrypoint to find global tables
    propFuncs = []
    for entry in eps:
        entryFunc = getFunctionAt(entry)
        propFuncs += EFIUtil.propogateFunctionVariables(entryFunc, callonExit=EFIUtil.findGlobalEfiPointers)

    # label instances of locate protocol
    EFIUtil.identifyGBSProtocol(propFuncs)

    # label instances of SMST
    EFIUtil.identifySMST()

    # label insances where SMST uses locate protocol
    EFIUtil.identifySMSTProtocol()

    # label Management Mode handlers
    EFIUtil.identifySMSTHandlers()

    # log what is found
    artifacts().logArtifacts()

    # attempt to find callouts
    EFIAnalytic.identifyCallouts()

    # attempt to find TOCTOU instances
    EFIAnalytic.identifyTOCTOU()

UEFISurveyor/gui_runner.py
dcg-red-team/UEFI_Surveyor
# Copyright 2023 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# @category UEFISurveyor.internal

from logger import logger
from uuid import UUID
import json


def convertGuidStr(name):
    # convert guid from memory
    ret = name[:8]
    ret += name[12:16]
    ret += name[8:12]
    ret += name[22:24]
    ret += name[20:22]
    ret += name[18:20]
    ret += name[16:18]
    ret += name[30:32]
    ret += name[28:30]
    ret += name[26:28]
    ret += name[24:26]
    return ret


def jsonKeystoUUID(inDict):
    outDict = {}
    for key in inDict:
        outDict[UUID(key)] = inDict[key]
    return outDict


class Guids(object):
    """Class to hold UEFI GUIDs"""
    def __init__(self):
        self.guidDict = {}

    def loadGuidFile(self, filename):
        # load UEFI Guid file
        with open(filename, 'r') as f:
            jsondata = f.read()
        self.guidDict.update(json.loads(jsondata, object_hook=jsonKeystoUUID))

    def getGuidName(self, guid):
        # Return name of GUID
        if guid in self.guidDict.keys():
            return self.guidDict[guid]
        return None

    def logGuids(self):
        # Log Guids extracted from guid file
        logger().log('found the following GUIDS')
        for key in self.guidDict.keys():
            logger.log()('{}, {}'.format(key, self.guidDict[key]))


_guids = Guids()


def guids():
    return _guids

UEFISurveyor/guids.py
dcg-red-team/UEFI_Surveyor
# Copyright 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# @category UEFISurveyor.internal

import logging
import sys
import time


class Logger:
    """Class for logging to console, text file"""
    def __init__(self, name=None, stream=False):
        """The Constructor assign streamhandler."""
        self.logfile = None
        self.logstream = None
        self.rootLogger = logging.getLogger('{}{}'.format(__name__, str(time.time())))
        self.rootLogger.setLevel(logging.INFO)
        if stream:
            self.enableStream
        if name:
            self.enableLogFile(name)

    def enableStream(self):
        if self.logstream:
            self.closeStream
        self.logstream = logging.StreamHandler(sys.stdout)
        self.rootLogger.addHandler(self.logstream)

    def closeStream(self):
        if self.logstream:
            self.rootLogger.removeHandler(self.logstream)
            self.logstream.flush()
            self.logstream = None

    def enableLogFile(self, name):
        """Sets the log file for the output."""
        # Close filehandler if enabled then set a new filehandler
        if self.logfile:
            self.closeLogFile()

        if name:
            # Open new log file and keep it opened
            try:
                # creates FileHandler for log file
                self.logfile = logging.FileHandler(filename=name, mode='w')
                self.rootLogger.addHandler(self.logfile)  # adds filehandler to root logger

            except Exception:
                print("WARNING: Could not open log file enabling console output'{}'".format(name))
                self.enableStream()

    def closeLogFile(self):
        """Closes the log file."""
        if self.logfile:
            try:
                self.rootLogger.removeHandler(self.logfile)
                self.logfile.close()
            except Exception:
                print("WARNING: Could not close log file")
            finally:
                self.logfile = None

    def log(self, text):
        """Log message to rootlogger"""
        self.rootLogger.log(logging.INFO, text)


_logger = Logger()


def logger():
    return _logger

UEFISurveyor/logger.py
dcg-red-team/UEFI_Surveyor
# Copyright 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import namedtuple
from multiprocessing import Process, SimpleQueue
import os
import subprocess
import platform

from options import getGhidraOptions, loadOptions
from UEFISurveyor.logger import logger


Job = namedtuple("Job", ['task', 'filename'])

END_TASK = 0
PROCESS_FILE = 1


class progressBar:
    def __init__(self, size):
        self.size = size
        self.completed = 0

    def update(self, something=None):
        self.completed += 1
        self.printProgress()

    def printProgress(self):
        size = 100
        x = int(size * (self.completed / self.size))
        print(f"Progress [{'#'*x}{'.'*(size-x)}] {x}% {self.completed}/{self.size}")

    def complete(self):
        return self.completed == self.size


def getFileDetails(filename):
    machine = None
    dataType = None
    machineType = None
    with open(filename, 'rb') as f:
        data = f.read(0x100)
        if b'VZ' in data:
            machine = data[2:4]
            dataType = 'TE'
        elif b'MZ' in data:
            dataType = 'PE'
            pe = data.find(b'PE')
            if pe > -1:
                machine = data[pe+4: pe+6]
    if machine:
        if machine == b'\x64\x86':
            machineType = '64bit'
        elif machine == b'\x4c\x01':
            machineType = '32bit'
    return (dataType, machineType)


def runHeadless(options, job, results, worker):
    ostype = platform.system()
    if ostype == 'Linux':
        gheadlesspath = os.path.join(options.ghidraPath, 'support', 'analyzeHeadless')
    elif ostype == 'Windows':
        gheadlesspath = os.path.join(options.ghidraPath, "support", "analyzeHeadless.bat")
    else:
        logger().log('Only supported on Windows and Linux at this time')
        return
    while True:
        injob = job.get()
        if injob.task == PROCESS_FILE:
            filename = injob.filename
            (dataType, machineType) = getFileDetails(filename)
            command = [gheadlesspath, options.ghidraProject, f'{options.projectName}{worker}', "-import", filename, "-overwrite"]
            command.append("-scriptPath")
            command.append(options.scriptPath)
            command.append("-preScript")
            command.append(os.path.join("UEFISurveyor", "analyze_pre.py"))
            if machineType == '64bit':
                command.append(options.gdtPath64)
            elif machineType == '32bit':
                command.append(options.gdtPath32)
            command.append("-postScript")
            command.append(os.path.join("UEFISurveyor", "analyze_post.py"))
            command.append(options.guidDBPath)
            command.append("-noanalysis")
            if dataType == 'TE':
                command.append('-processor')
                if machineType == '64bit':
                    command.append('x86:LE:64:default')
                elif machineType == '32bit':
                    command.append('x86:LE:32:default')
            command_txt = ' '.join(i for i in command)
            logger().log("About to execute {}".format(command_txt))
            res = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

            if 'Import failed for file' in str(res):
                logger().log(f'Unable to process {filename}')

            results.put(res)
        elif injob.task == END_TASK:
            break
        else:
            logger().log('Problem with the queue', injob.task)


if __name__ == "__main__":
    logger().enableStream()
    logger().log("Starting UEFI Surveyor")
    alloptions = loadOptions(os.path.join('UEFISurveyor', 'options.yaml'))
    options = getGhidraOptions(alloptions)
    proc_count = 3
    work = SimpleQueue()
    result = SimpleQueue()
    procs = []

    for i in range(proc_count):
        p = Process(target=runHeadless, args=(options, work, result, i))
        procs.append(p)
    for p in procs:
        p.start()

    pb = progressBar(len(os.listdir(options.efiPath)))
    pb.printProgress()

    filelist = os.listdir(options.efiPath)
    count = 0

    for _ in range(proc_count):
        if filelist:
            work.put(Job(PROCESS_FILE, os.path.join(options.efiPath, filelist.pop())))
        else:
            break
    while filelist:
        for count in range(proc_count):
            if filelist:
                work.put(Job(PROCESS_FILE, os.path.join(options.efiPath, filelist.pop())))
            else:
                break
        if count:
            for _ in range(count + 1):
                res = result.get()
                pb.update()

    while not pb.complete():
        res = result.get()
        pb.update()

    for p in procs:
        work.put(Job(END_TASK, None))

    for proc in procs:
        proc.join()

    logger().log("fin")

headless_runner.py
dcg-red-team/UEFI_Surveyor
# Copyright 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# @category UEFISurveyor.headless

from collections import namedtuple
import yaml


surveyorOptions = namedtuple("UEFISurveyor_Options", ['ghidraPath', 'ghidraProject', 'scriptPath', 'guidDBPath', 'gdtPath64', 'gdtPath32', 'projectName', 'efiPath'])


def loadOptions(filename='options.yaml') -> dict:
    try:
        with open(filename, 'r') as f:
            data = f.read()
    except IOError:
        data = ''
    try:
        options = yaml.safe_load(data)
    except yaml.YAMLError:
        options = {}
    return options


def getGhidraOptions(options: dict) -> surveyorOptions:
    ghidrapath = options['Analysis']['Ghidrapath']
    ghidraprojects = options['Analysis']['Projectpath']
    scriptpath = options['Analysis']['Scriptpath']
    guiddbpath = options['Analysis']['GUIDDB']
    gdtpath64 = options['Analysis']['GDT64']
    gdtpath32 = options['Analysis']['GDT32']
    projectname = options['Analysis']['Projectname']
    efipath = options['Analysis']['EFIPath']
    retOptions = surveyorOptions(ghidrapath, ghidraprojects, scriptpath, guiddbpath, gdtpath64, gdtpath32, projectname, efipath)
    return retOptions

options.py
dcg-red-team/UEFI_Surveyor