pwndbg/scripts/_docs/build_function_docs.py

"""
If the PWNDBG_DOCGEN_VERIFY environment variable
is set, then    : Exit with non-zero exit status if the docs/functions/ files
                  aren't up to date with the sources. Don't modify anything.

If it isn't, this fixes up the docs/functions/ files to be up
to date with the information from the sources.
"""

from __future__ import annotations

import json
import os
import sys
from typing import Dict
from typing import Tuple

from mdutils.mdutils import MdUtils

from scripts._docs.function_docs_common import BASE_PATH
from scripts._docs.function_docs_common import ExtractedFunction
from scripts._docs.function_docs_common import extracted_filename
from scripts._docs.gen_docs_generic import ALL_DEBUGGERS
from scripts._docs.gen_docs_generic import update_files_simple
from scripts._docs.gen_docs_generic import verify_existence
from scripts._docs.gen_docs_generic import verify_files_simple

INTRO_TEXT = """
pwndbg provides a set of functions which can be used during expression evaluation to
quickly perform common calculations. These can even be passed to other commands as arguments.
Currently, they only work in gdb.

To see a list of all functions, including those built into gdb, use `help function`. To see
the help of any given function use `help function function_name`. Function invokation must
include a preceding $ sign and must include brackets. For instance, invoke the `environ`
function like so:
```
pwndbg> p $environ("LANG")
$2 = (signed char *) 0x7fffffffe6da "LANG=en_US.UTF-8"
```
If the result of the function is being passed to a pwndbg command, make sure to either escape
the function argument's quotes, or put the whole function call in quotes.
```
pwndbg> tele $environ("LANG")
usage: telescope [-h] [-r] [-f] [-i] [address] [count]
telescope: error: argument address: debugger couldn't resolve argument '$environ(LANG)':
    No symbol "LANG" in current context.
pwndbg> tele $environ(\\"LANG\\")
00:0000│  0x7fffffffe6cf ◂— 'LANG=en_US.UTF-8'
01:0008│  0x7fffffffe6d7 ◂— 'US.UTF-8'
02:0010│  0x7fffffffe6df ◂— 0x4e49475542454400
[...]
pwndbg> tele '$environ("LANG")'
00:0000│  0x7fffffffe6cf ◂— 'LANG=en_US.UTF-8'
01:0008│  0x7fffffffe6d7 ◂— 'US.UTF-8'
02:0010│  0x7fffffffe6df ◂— 0x4e49475542454400
[...]
```
## pwndbg functions
""".strip()


def convert_to_markdown(extracted: list[Tuple[str, list[ExtractedFunction]]]) -> Dict[str, str]:
    """
    Returns:
        A dict which maps filenames to their markdown contents.
        It will have only one item (the index.md).
    """
    markdowned = {}

    mdFile = MdUtils(INDEX_PATH)
    mdFile.new_header(level=1, title="Functions")
    mdFile.new_paragraph(INTRO_TEXT)

    all_functions: set[str] = set()
    for _, funcs in extracted:
        all_functions.update([func.name for func in funcs])

    for func_name in sorted(all_functions):
        # Make a list of all debuggers that know about this
        # function. Make sure they all agree on the contents.
        supported_debuggers: list[str] = []
        func: ExtractedFunction | None = None

        for debugger, dfuncs in extracted:
            # Slow but whatever
            for dfunc in dfuncs:
                if func_name == dfunc.name:
                    supported_debuggers.append(debugger)
                    if not func:
                        func = dfunc
                    elif func != dfunc:
                        print(f"Error: Debuggers don't agree on {func.name}")
                        print(f"{supported_debuggers[0]} says: {func}\n")
                        print(f"{debugger} says: {dfunc}")
                        exit(10)

        mdFile.new_paragraph(f"### **{func.name}**")
        func_signature_code = f"""
``` {{.python .no-copy}}
{func.name}{func.signature}
```
"""
        if (
            " object at " in func.signature or "<" in func.signature
        ):  # '>' is valid in type annotation (->)
            print(f'Signature of {func.name} is rendered as "{func.signature}", please edit')
            print("the sanitize_signature() function (in the extractor) to display the ")
            print("signature better in the docs.")
            sys.exit(5)

        mdFile.new_paragraph(func_signature_code)
        mdFile.new_paragraph(func.docstring.replace("Example:", "#### Example"))
        mdFile.new_paragraph("-" * 10)

    hide_nav = "---\nhide:\n  - navigation\n---\n"
    autogen_warning = (
        "<!-- THIS WHOLE FILE IS AUTOGENERATED. DO NOT MODIFY IT. See scripts/generate-docs.sh -->"
    )
    markdowned[INDEX_PATH] = hide_nav + autogen_warning + "\n" + mdFile.get_md_text()
    return markdowned


def read_extracted() -> list[Tuple[str, list[ExtractedFunction]]]:
    """
    Read json files from disk.

    Returns:
        A list of tuples of the form: (debugger name, list of
        convenience functions that debugger supports).
    """

    result: list[Tuple[str, list[ExtractedFunction]]] = []

    for debugger in ALL_DEBUGGERS:
        filepath = extracted_filename(debugger)
        print(f"Consuming {filepath}..")

        with open(filepath, "r") as file:
            raw_data = json.loads(file.read())

        # Convert the dict objs to ExtractedFunction
        data = [ExtractedFunction(**func) for func in raw_data]

        result.append((debugger, data))

        # We consumed the temporary file, we can delete it now.
        os.remove(filepath)

    return result


INDEX_PATH = os.path.join(BASE_PATH, "index.md")


def main():
    if len(sys.argv) > 1:
        print("This script doesn't accept any arguments.")
        print("See top of the file for usage.")
        sys.exit(1)

    just_verify = False
    if os.getenv("PWNDBG_DOCGEN_VERIFY"):
        just_verify = True

    print("\n==== Function Documentation ====")

    extracted = read_extracted()
    markdowned = convert_to_markdown(extracted)
    assert len(markdowned) == 1  # Only index.md

    if just_verify:
        print("Checking if all files are in place..")
        missing, extra = verify_existence(list(markdowned.keys()), BASE_PATH)
        if missing or extra:
            print("To add mising files please run ./scripts/generate-docs.sh.")
            print("To remove extra files please remove them manually.")
            sys.exit(2)
        print("Every file is where it should be!")

        print("Verifying contents...")
        err = verify_files_simple(markdowned)
        if err:
            print("VERIFICATION FAILED. The files differ from what would be auto-generated.")
            print("Error:", err)
            print("Please run ./scripts/generate-docs.sh from project root and commit the changes.")
            sys.exit(3)

        print("Verification successful!")
    else:
        print("Updating files...")
        update_files_simple(markdowned)
        print("Update successful.")

        missing, extra = verify_existence(list(markdowned.keys()), BASE_PATH)
        assert (
            not missing
            and "Some files (and not the index) are missing, which should be impossible."
        )

        if extra:
            print("Please delete the extra files by hand.")
            sys.exit(4)


if __name__ == "__main__":
    main()