workname: Делаем директории в IDA и автоматизируем это.
# Пишем простой скрипт: Распределение функций C++ по директориям в IDA.

Задача: Есть проект, где в каждом модуле есть куча классов с сопутствующими функциями, которые не организованы. Нужно написать скрипт, который парсит функции и распределяет их по директориям (новая фича IDA 7.7).

Цели:
- Распарсить функции
- Применить демангл
- Организовать директории по классам

Ссылки:
- AlphaGo: [categorize_go_folders.py](https://github.com/SentineLabs/AlphaGolang/blob/main/3.categorize_go_folders.py)
- Док-я IDA (SDKDOC) #1 : [Ссылка](https://hex-rays.com/products/ida/support/sdkdoc/dirtree_8hpp.html)
- Док-я IDA (IDAPYTHONDOC) #2: [Ссылка](https://www.hex-rays.com/products/ida/support/idapython_docs/ida_dirtree.html)

# Заметки

- Функции в демангл-версии выглядят как `ModuleName::FuncName(arg)`
- Функция для получения списка функций в IDB: `idautils.Functions()`
- Функция для получения имени функции из объекта функции: `idc.get_func_name(func)`
- Функция для демангла имени: `string demangle_name(string name, long disable_mask);`
- `ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS)` - судя по всему корень дерева директорий, по умолчанию функции кладутся туда. (root node/корневой узел-директория)
- dirtree_t.resolve_path(dirname:str) - возвращает direntry_t. из нее можно вытащить .idx - номер директории
- get_entry_name(idx) - наоборот, возвращает только имя директории
- idx = 0 - походу это всегда './'
- root_node.get_dir_size(0) - вернуть размер корня.

# Проблемы

- Функции, сгенерированный по шаблонам. Имеют чуть другой mangling алгоритм. **UPD: не проблема, в имя модуля уходит весь шаблон.**
- Останется ряд функций вида sub_XXXXXXXX, которые, частично, составляют vtable различных модулей. Мб попробовать их как-то научится парсить?
- Не все функции деманглятся. Хз почему, наверное алгоритм в компиляторе изначально чуть другой?
- Некоторые названия модулей образуются с помощью префикса/суффикса. Можно ли как-то мерджить директории?
- Если переименовать функцию, то она по умолчанию останется в той же директории, что и была. TODO: при старте получать список existed директорий и удалять их все. И соответственно снова наполнять с уже обновленными именами.
- апхахпахпахп я не знаю почему но findnext возвращает False раньше чем закончатся файлы в директории. **UPD** Это происходило из-за того что я перемещал файл сразу после того, как находил. Это почему-то ломало курсор кекв.

TODO: Пофиксить ошибку, если в функе нет неймспеса, а в параме есть ( verify(std::string) -> verify(std )

# Фичи для реализации

- Сортировка по префиксу-подстроке
- Сделать Favorite по ПКМ.
- Сделайть отдельный GUI
- Сделать Favorite-Auto по регексу

- Сделать для Rust
- Сделать для Golang
# Алгоритм
- Чтобы увидеть директории: Functions -> ПКМ -> Show Folders
- Получаем список функций при помощи `for func in idautils.Functions():`
- Они будут иметь 'оригинальное' имя вида `AfxThrowMemoryException@@YGXXZ`
- Парсим имя при помощи `demangle_name(str, disable_mask)`
- str - имя функции
- disable_mask - режим по которому раскукоживается название функции. Имеет два режима:
- get_inf_attr(INF_LONG_DN). Раскукожить полностью, т.е получить полную сигнатуру функции. Имеем на выходе названия вида:
- `public: virtual unsigned int __thiscall CFile::Read(void *, unsigned int)`
- get_inf_attr(INF_SHORT_DN). Раскукожить только модуль, имя и аргументы. Имеем на выходе:
- `CFile::Read(void *,uint)`
- В наших целях нам будет достаточно только 'короткого' имени функции. Сплитаем по `::` и получаем имя модуля через [0].
- WARNING: это некорректный алгоритм, так как прегенерированные шаблонные функции имеют чуть другой demangled name. Надо обработать эту ветку.
- По результату демангла функции можно поделить на следующие категории:
- Модульные - те, которые имеют :: в названии (?)
- Сишные - те, которые деманглятся, но имеют обычное сишное имя (без ::)
- ~~У которых произошла ошибка с деманглом и остальные символы (требуют ручного переноса в директории или изменить имя функции вручаную)~~
- Unrecognized - стандартные sub_XXXXXXXX.
- Для работы с директориями функций в иду добавили следующее API:

# Алгоритм Создания/Удаления новых нод

# Алгоритм Обхода нод дерева.

# Алгоритм поиска директорий
Все пути можно получить в динамике. Но почитав документацию можно предположить, что должен быть механизм поиска потомков ноды dirent_t файловой системы
# Готовые скрипты
Простой скрипт поиска и переименования

import idautils
import ida_dirtree
import idc 

print("="*20)
modules = {}

for fun in idautils.Functions():
    name = idc.get_func_name(fun)
    
    name_dmngl_long  = demangle_name(name, get_inf_attr(INF_LONG_DN))
    name_dmngl_short = demangle_name(name, get_inf_attr(INF_SHORT_DN))
    
    if name_dmngl_short:
        if "::" in name_dmngl_short:
            module_name = '[' + name_dmngl_short.split("::")[0] + ']'
            modules.setdefault(module_name, []).append(fun)
        else:
            modules.setdefault("Unrecognized", []).append(fun)
    else: 
        if name.startswith("sub_"):
            modules.setdefault("sub_", []).append(fun)
        elif '::' in name:
            module_name = '[' + name.split("::")[0] + ']'
            modules.setdefault(module_name, []).append(fun)
        else:
            modules.setdefault("Other", []).append(fun)
            
for module in modules.keys():
    ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS).mkdir(module)
    print('[D] Created folder for ', module)
    for fun in modules[module]:
        name = idc.get_func_name(fun)
        ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS).rename(name, module+"/"+name)

print("="*20)

IN_DEVELOPMENT:

import idautils
import ida_dirtree
import idc 

#######################################################################################
# doc = https://www.hex-rays.com/products/ida/support/idapython_docs/ida_dirtree.html #
#######################################################################################

print("="*20)
modules = {}

for fun in idautils.Functions():
    name = idc.get_func_name(fun)
    
    name_dmngl_long  = demangle_name(name, get_inf_attr(INF_LONG_DN))
    name_dmngl_short = demangle_name(name, get_inf_attr(INF_SHORT_DN))
    
    if name_dmngl_short:
        if "::" in name_dmngl_short:
            module_name = '[' + name_dmngl_short.split("::")[0] + ']'
            modules.setdefault(module_name, []).append(fun)
        else:
            modules.setdefault("Unrecognized", []).append(fun)
    else: 
        if name.startswith("sub_"):
            modules.setdefault("sub_", []).append(fun)
        elif '::' in name:
            module_name = '[' + name.split("::")[0] + ']'
            #print(module_name, name)
            modules.setdefault(module_name, []).append(fun)
        else:
            modules.setdefault("Other", []).append(fun)

# delete previosly created module dirs

# 1. Получаем direntry_t на нашу директорию
a = ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS).resolve_path('[std]')
print(a.idx)
# 2. Делаем на основе нее курсор для обхода, пихаем в idx значение нашей direntry_t
dt = ida_dirtree.dirtree_cursor_t()
dt.parent = a.idx
print(ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS).get_entry_name(a)) # self-check
# 3. Делаем итератор на основании курсора.
ite = ida_dirtree.dirtree_iterator_t()
ite.cursor = dt
# 4. Обходим при помощи findnext и получаем полный путь при помощи get_abspath(iter)
print(ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS).findnext(ite))
print(ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS).get_abspath(ite.cursor))
print(ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS).findnext(ite))
print(ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS).get_abspath(ite.cursor))

# Make folders with funcs
for module in modules.keys():
    ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS).mkdir(module)
    #print('[D] Created folder for', module, "\n    Functions:", len(modules[module]))
    for fun in modules[module]:
        name = idc.get_func_name(fun)
        #print("\t", name)
        ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS).rename(name, module+"/"+name)

print("="*20)

import ida_dirtree
import idautils
import idc
print('-'*10)

func_dir_root = ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS)

def harvest_modules_names():
    modules_list = []
    for func in Functions():
        name = idc.get_func_name(func)
        demangled_short = demangle_name(name, get_inf_attr(INF_SHORT_DN))
        
        if demangled_short:
            if '::' in demangled_short:
                nm = demangled_short.split('::')[0]                      
                modules_list.append(nm)
    return set(modules_list)

ml = harvest_modules_names()
for i in ml:
    module_idx= func_dir_root.resolve_path(i).idx
    if module_idx != 0xFFFFFFFF:
        print(i, module_idx)

Function window sorting for C++:

import ida_dirtree
import idc

import re
print('-'*20)

root_node: ida_dirtree.dirtree_t
root_node = ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS) # root tree instance for Functions window

# (Folder Attribute Label) Append number of functions in module to module folder name
FLAG_ADD_COUNT_TO_MODULE_NAME = True
# Align module folder name. For readable attributes labels.
FLAG_ALIGNED_MODULE_NAMES     = True
# Create prefix for matched names of folders. Like Favorite and for sorting. Hack: Use whitespace symbols.
# TODO: FIX: Can be nested match that will create nested prefix. Like [C... and [CFG...
FLAG_CURSED_FAVORITE          = True
CURSED_PREFIX_MATCH = {
    "\[Rt[\d\w]+\]" : "\u2606",
    "C[A-Z]{1}[a-z]{1}[\d\w]+"  : "\u2730",
    "\[RtT[\d\w]+<[\d\w,\*\s]+>\]"  : "\u2726"
    
}

def create_functions_modules_dict() -> dict:
    '''
    Create dictionary with modules as keys and functions names as values.
    Return - dict.
    '''
    modules = {}
    
    #iterate over functions in IDB
    for func in Functions():
        #get function name from IDB
        name = idc.get_func_name(func)
        #it will not be demangled. demangle it:
        dmngl_name = demangle_name(name, get_inf_attr(INF_SHORT_DN)) #it return NULL on error
        
        module_name = ""
        if dmngl_name:
            # we have smth that looks like C/C++ name mangling.
            # it can have ::  - namespace delimeter. Split by it and get module name:
            if '::' in dmngl_name:
                module_name = '[{}]'.format(dmngl_name.split('::')[0])
            
            # if not - it some generic function (or without module namespace)
            else:
                module_name = "\u3000 Unrecognized"
        else:
            # we have some error with demangling...
            # ... because sometimes IDA put already demangled name to IDB¯\_(ツ)_/¯
            # ... or with functions already renamed by hand.
            if '::' in name:
                module_name = '[{}]'.format(name.split('::')[0])
                
            # ... with standard unrecognized funcs 
            elif name.startswith("sub_"):
                module_name = "\u3000 sub_"
                
            # ... when it smth other. "_start" as example
            else:
                module_name = "\u3000 Unrecognized"
            
        # and add original name to specific's module list
        modules.setdefault(module_name, []).append(name)
        
    return modules

def generate_functions_folders(modules: dict):
    '''
    Generate folders in Function window with ida_dirtree API
    based on modules dictionary (use create_modules_dict()).
    '''
    format_for_attrs = '{}'
    
    # Check comments at start of file for flags explanation
    if FLAG_ALIGNED_MODULE_NAMES:
        max_len_of_module_name = max([len(i) for i in modules.keys()])
        format_for_attrs = "{{:<{}}}".format(max_len_of_module_name+2)     
    
    print("Generated folders (Functions): [Module Namespace] [number of f:unctions]")
    
    #iterate over module names
    for module in modules.keys():
        module_folder = format_for_attrs.format(module)
        
        # Check comments at start of file for flags explanation
        if FLAG_ADD_COUNT_TO_MODULE_NAME: 
            module_folder += "[f: {}]".format(len(modules[module]))
        
        # Check comments at start of file for flags explanation
        if FLAG_CURSED_FAVORITE:
            for pref in CURSED_PREFIX_MATCH.keys():
                #if module.startswith(pref):
                if re.match(pref, module):
                    module_folder = CURSED_PREFIX_MATCH[pref] + module_folder
        
        print(module_folder)
        # using ida_dirtree api for creating folder at root of Functions window
        err = root_node.mkdir(module_folder)
        if err:
            print(err, module, root_node.errstr(err))
            continue
            
        for function in modules[module]:
            # little cursed ida_dirtree API: moving files in folders with "tree.rename"
            err = root_node.rename(function, module_folder+'/'+function)
            if err:
                print(err, function, root_node.errstr(err))
                continue

def delist_functions_folders(start_path):
    '''
    Delete all folders at start_path in Functions windows and move functions to root directory node.
    CTRL+Z works, if you need to revert changes.
    '''
    # create standard iterator
    node_iter = ida_dirtree.dirtree_iterator_t()
    # find first node in direntry_t.      V - pattern used for search, so start_path also prefix for nested folders
    ok = root_node.findfirst(node_iter, start_path+'/*')
    paths = []
    dirs  = []
    while ok:
        # getting absolute path for node. direntry_t doesn't have any "get childs" functions?
        dirent_path = root_node.get_abspath(node_iter.cursor)
        
        if root_node.isfile(dirent_path):
            paths.append(dirent_path) # if node is file - adding it to files list
        elif root_node.isdir(dirent_path):
            dirs.append(dirent_path)  # if node is dir - recursive fall to this dir
            delist_functions_folders(dirent_path) # and add for deleting list in future
        else:
            print("wtf")    # never must be branched 
            raise Exception('wtf')
        ok = root_node.findnext(node_iter)
        
    for p in paths:
        root_path = '/'+p.split('/')[-1]
        err = root_node.rename(p, root_path)  # move all found files to root directory
        if err:
            print(err, root_node.errstr(err))
            
    for d in dirs:
        err = root_node.rmdir(d)              # and delete all folders
        if err:
            print(err, root_node.errstr(err))
            
delist_functions_folders('/')    
modules = create_functions_modules_dict()
generate_functions_folders( modules )