# Import the results of the cwe_checker as bookmarks and comments into Ghidra.
#
# Usage:
# - Copy this file into the Ghidra scripts folder
# - Run the cwe_checker on a binary and save its output as a json file, e.g. with
#   "cwe_checker BINARY --json --out output.json"
# - Open the binary in Ghidra and run this file as a script. Select the generated json file when prompted.

import json


def bookmark_cwe(ghidra_address, text):
    previous_bookmarks = getBookmarks(ghidra_address)
    for bookmark in previous_bookmarks:
        if '[cwe_checker]' == bookmark.getCategory():
            if text not in bookmark.getComment():
                createBookmark(ghidra_address, '[cwe_checker]', bookmark.getComment() + '\n' + text)
            return
    createBookmark(ghidra_address, '[cwe_checker]', text)
    return


def comment_cwe_eol(ghidra_address, text):
    old_comment = getEOLComment(ghidra_address)
    if old_comment is None:
        setEOLComment(ghidra_address, text)
    elif text not in old_comment:
        setEOLComment(ghidra_address, old_comment + '\n' + text)


def comment_cwe_pre(ghidra_address, text):
    old_comment = getPreComment(ghidra_address)
    if old_comment is None:
        setPreComment(ghidra_address, text)
    elif text not in old_comment:
        setPreComment(ghidra_address, old_comment + '\n' + text)


def get_cwe_checker_output():
    ghidra_file = askFile('Select json output file of the cwe_checker', 'Open')
    with open(ghidra_file.getAbsolutePath()) as json_file:
        return json.load(json_file)


def main():
    """
    Annotate cwe_checker results in Ghidra as end-of-line
    comments and bookmarks to the corresponding addresses.
    """
    warnings = get_cwe_checker_output()
    for warning in warnings:
        if len(warning['addresses']) == 0:
            cwe_text =  '[' + warning['name'] + '] ' + warning['description']
            ghidra_address = currentProgram.getMinAddress().add(0)
            bookmark_cwe(ghidra_address, cwe_text)
            comment_cwe_pre(ghidra_address, cwe_text)
        else:
            address_string = warning['addresses'][0]
            ghidra_address = currentProgram.getAddressFactory().getAddress(address_string)
            bookmark_cwe(ghidra_address, warning['description'])
            comment_cwe_eol(ghidra_address, warning['description'])

main()