pwndbg/scripts/_gen_function_docs.py

#!/usr/bin/env python
"""
You should use scripts/generate_docs.sh and scripts/verify_docs.sh instead
of using this.

If the PWNDBG_GEN_DOC_JUST_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 os
import re
import sys
from inspect import getdoc
from inspect import signature
from typing import Dict

from mdutils.mdutils import MdUtils

import pwndbg
from pwndbg.gdblib.functions import _GdbFunction
from scripts._gen_docs_generic import update_files_simple
from scripts._gen_docs_generic import verify_existence
from scripts._gen_docs_generic import verify_files_simple


def extract_functions() -> Dict[str, _GdbFunction]:
    """
    Returns a dictionary that mapes function names to
    the corresponding _GdbFunction objects.
    """
    functions = pwndbg.gdblib.functions.functions
    result = {}

    for f in functions:
        result[f.name] = f

    return result


def sanitize_signature(func_name: str, sig: str) -> str:
    """
    We need to strip ' from type annotations, and cleanup
    some functions that don't display properly.
    """
    sig = sig.replace("'", "")
    match func_name:
        case "fsbase":
            sig = re.sub(r"<gdb\.Value object at 0x[0-9a-fA-F]+>", "gdb.Value(0)", sig)
        case "gsbase":
            sig = re.sub(r"<gdb\.Value object at 0x[0-9a-fA-F]+>", "gdb.Value(0)", sig)
    return sig


def convert_to_markdown(named_funcs: Dict[str, _GdbFunction]) -> 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")
    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
    """
    mdFile.new_paragraph(intro_text.strip())

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

        mdFile.new_paragraph(func_signature_code)
        mdFile.new_paragraph(
            "#### Description\n" + getdoc(func).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 check_index(scoped_params: Dict[str, list[Parameter]]):
    assert (
        len(scoped_params.keys()) == 3
        and "It seems a new scope has been added, "
        f"please update the index file ({index_path}) and bump this number accordingly."
    )


base_path = "docs/functions/"  # Must have trailing slash.
index_path = base_path + "index.md"

# ==== Start ====

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_GEN_DOC_JUST_VERIFY"):
    just_verify = True

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

named_functions = extract_functions()
markdowned = convert_to_markdown(named_functions)
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)