cthaeh/apply_dta.py at main · jeffaf/cthaeh · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
# apply_dta.py - Ghidra Pre-Script: Load Talos Windows Driver Data Type Archive
#
# Loads the Cisco Talos .gdt archive into the current program's data type manager
# so that Windows kernel types (IRP, IO_STACK_LOCATION, DEVICE_OBJECT, etc.) are
# available during analysis.
#
# Usage: analyzeHeadless ... -preScript apply_dta.py
#
# Source: https://github.com/Cisco-Talos/Windows-drivers-GDT-file
#
# NOTE: This is Jython (Python 2). No f-strings. Use % formatting.
#
# @category Security
# @author Jeff Barron

import os

from ghidra.program.model.data import FileDataTypeManager
from java.io import File


def find_gdt():
    """Find the .gdt file relative to this script's location."""
    # Try relative to script directory first
    script_dir = os.path.dirname(sourceFile.getAbsolutePath())
    candidates = [
        os.path.join(script_dir, "data", "windows_driver_types.gdt"),
        os.path.join(script_dir, "windows_driver_types.gdt"),
    ]

    # Also check CTHAEH_DTA_PATH env var
    env_path = os.environ.get("CTHAEH_DTA_PATH", "")
    if env_path:
        candidates.insert(0, env_path)

    for path in candidates:
        if os.path.isfile(path):
            return path

    return None


def apply_dta():
    """Load the .gdt archive into the current program's data type manager."""
    gdt_path = find_gdt()

    if gdt_path is None:
        println("[apply_dta] WARNING: No .gdt file found. Skipping DTA import.")
        println("[apply_dta]   Run: python download_dta.py")
        println("[apply_dta]   Or set CTHAEH_DTA_PATH environment variable.")
        return

    println("[apply_dta] Loading DTA: %s" % gdt_path)

    gdt_file = File(gdt_path)
    archive = FileDataTypeManager.openFileArchive(gdt_file, False)

    try:
        dtm = currentProgram.getDataTypeManager()
        source_dtm = archive.getDataTypeManager()

        count = 0
        txn = dtm.startTransaction("Import Talos DTA")
        try:
            for dt in source_dtm.getAllDataTypes():
                name = dt.getName()
                # Skip built-in types that Ghidra already has
                path = dt.getCategoryPath().getPath()
                if path.startswith("/BuiltInTypes"):
                    continue
                dtm.addDataType(dt, None)
                count += 1
        finally:
            dtm.endTransaction(txn, True)

        println("[apply_dta] Imported %d data types from Talos archive." % count)

    finally:
        archive.close()


# Entry point
apply_dta()